Skip to content

Derivatives

This example takes derivatives with ppy.grad and ppy.value_and_grad, which follow one rule table on every path, so the nine digits this program prints are the same nine everywhere. ppy.grad(f) is the gradient of f with respect to its first parameter. ppy.value_and_grad(f, argnums=1) is the value and the gradient with respect to another.

Run it

python  gradients.ppy
ppy run gradients.ppy

What it prints

python gradients.ppy, ppy run gradients.ppy

2.929684375 4.201088436
-2.488585914 0.523598776 4.7e-17

Reverse mode on both paths

def f(x: float, y: float) -> float:
    return math.sin(x) * y + x * x


df = ppy.grad(f)
both = ppy.value_and_grad(f, argnums=1)

Under CPython the derivative is made from f's source the first time it is called: the body restated one operation at a time, then every operation's adjoint in reverse.

Natively the autodiff transform does the same to the IR, reverse mode over the core and tensor dialects, and slope(x, y) compiles to a call to the derived function. There is no finite differencing and no second implementation to drift.

A derivative is an ordinary function

def newton(x: float) -> float:
    for _ in range(6):
        x = x - g(x) / dg(x)
    return x

dg is used inside a loop like any callable. newton is native, and the root it finds is a root: the last column prints g at that root, zero to floating-point precision.

Compared with JAX and PyTorch

Newton's method on g from 100,000 starting points near 0.4, six steps each, with the derivative from each framework's own autodiff. The programs are in compare/: gradients_bench.ppy, gradients_jax.py, gradients_torch.py. Milliseconds for the whole batch, best of five, over five processes.

PPy: ppy.grad(g) is a function. newton calls it in a loop, and a loop over the starting points calls newton. Everything is scalar and native; nothing is batched:

dg = ppy.grad(g)


def newton(x: float) -> float:
    for _ in range(6):
        x = x - g(x) / dg(x)
    return x


def newton_all(count: int) -> float:
    total = 0.0
    for i in range(count):
        total += newton(0.4 + i * 1e-6)
    return total / count

JAX: jax.grad over jnp functions. To run 100,000 solves it wants them batched: vmap under jit, the loop as lax.fori_loop so it traces, and jax_enable_x64 so the digits match:

def newton(x):
    def step(_, x):
        return x - g(x) / dg(x)

    return jax.lax.fori_loop(0, 6, step, x)


newton_all = jax.jit(jax.vmap(newton))

PyTorch: torch.func.grad and vmap. The six steps stay a Python loop over a 100,000-element tensor in float64. The thread count is pinned to the performance cores', since its default of one per logical core spins on this hybrid CPU and the batch takes hundreds of milliseconds some runs:

dg = grad(g)


def newton(x):
    for _ in range(6):
        x = x - g(x) / dg(x)
    return x


newton_all = vmap(newton)
PPy JAX vmap + jit PyTorch torch.func
newton, 100k starts 11.89 ± 0.07 1.23 ± 0.10 6.52 ± 1.04

The derivative is the same nine digits in all three: reverse mode over the same rule table. What differs is the shape of the program.

JAX and PyTorch are fast here because the problem batches, and XLA and ATen vectorize across the batch. A scalar newton in a Python loop would cost them tens of microseconds per call. PPy's scalar loop pays nothing per call and is not vectorized across the batch.

Which one is faster depends on whether your problem comes as 100,000 independent solves or as one.

Intel Core Ultra 9 386H; JAX 0.11.1 and PyTorch 2.14.0 (CPU) on CPython 3.13.13, PPy on CPython 3.13.13, from a checkout on a native filesystem.

Limitations

What cannot be differentiated is refused:

  • a branch or a loop inside f is E1660
  • a non-float signature is E1661
  • an effect no derivative follows (I/O, a write, a thread) is E1662

Read on: Derivatives · The IR

gradients.ppy is hand-written; there is no .py source and no conversion step.

36_autodiff/gradients.ppy

import math

import ppy


def f(x: float, y: float) -> float:
    return math.sin(x) * y + x * x


def g(x: float) -> float:
    return math.exp(-x * x) * math.cos(3.0 * x)


df = ppy.grad(f)
both = ppy.value_and_grad(f, argnums=1)
dg = ppy.grad(g)


def slope(x: float, y: float) -> float:
    return df(x, y)


def pair(x: float, y: float) -> float:
    v, grad_y = both(x, y)
    return v * 2.0 + grad_y


def newton(x: float) -> float:
    for _ in range(6):
        x = x - g(x) / dg(x)
    return x


def main() -> None:
    print(f"{slope(0.7, 2.0):.9f} {pair(0.7, 2.0):.9f}")
    print(f"{dg(0.3):.9f} {newton(0.4):.9f} {g(newton(0.4)):.1e}")


main()

Counterpart programs

The programs the comparison above measured, each written the way its tool expects. The PPy one is first.

gradients_bench.ppy (PPy) 
"""The derivatives of `gradients.ppy`, and Newton's method from 100,000 starts: the PPY side."""

import math
import time

import ppy


def f(x: float, y: float) -> float:
    return math.sin(x) * y + x * x


def g(x: float) -> float:
    return math.exp(-x * x) * math.cos(3.0 * x)


df = ppy.grad(f)
both = ppy.value_and_grad(f, argnums=1)
dg = ppy.grad(g)


def slope(x: float, y: float) -> float:
    return df(x, y)


def pair(x: float, y: float) -> float:
    v, grad_y = both(x, y)
    return v * 2.0 + grad_y


def newton(x: float) -> float:
    for _ in range(6):
        x = x - g(x) / dg(x)
    return x


def newton_all(count: int) -> float:
    """The mean root over `count` starting points near 0.4, one Newton solve each."""
    total = 0.0
    for i in range(count):
        total += newton(0.4 + i * 1e-6)
    return total / count


def worst_residual(count: int) -> float:
    worst = 0.0
    for i in range(count):
        worst = max(worst, abs(g(newton(0.4 + i * 1e-6))))
    return worst


def main() -> None:
    print(f"{slope(0.7, 2.0):.9f} {pair(0.7, 2.0):.9f}")
    newton_all(1000)
    best = 1e9
    mean = 0.0
    for _ in range(5):
        started = time.perf_counter()
        mean = newton_all(100_000)
        best = min(best, (time.perf_counter() - started) * 1000.0)
    print(f"# newton, 100k starts: {best:.2f} ms")
    print(f"{mean:.9f} {worst_residual(100_000):.0e}")


main()
gradients_jax.py (Python) 
"""The same derivatives under JAX: `jax.grad`, and `vmap` over the starting points, jitted."""

import time

import jax
import jax.numpy as jnp

jax.config.update("jax_enable_x64", True)


def f(x, y):
    return jnp.sin(x) * y + x * x


def g(x):
    return jnp.exp(-x * x) * jnp.cos(3.0 * x)


df = jax.grad(f)
both = jax.value_and_grad(f, argnums=1)
dg = jax.grad(g)


def newton(x):
    def step(_, x):
        return x - g(x) / dg(x)

    return jax.lax.fori_loop(0, 6, step, x)


newton_all = jax.jit(jax.vmap(newton))


def timed(label, run):
    best = 1e9
    answer = None
    for _ in range(5):
        started = time.perf_counter()
        answer = run()
        best = min(best, time.perf_counter() - started)
    print(f"# {label}: {best * 1000:.2f} ms")
    return answer


def main():
    v, grad_y = both(0.7, 2.0)
    print(f"{float(df(0.7, 2.0)):.9f} {float(v * 2.0 + grad_y):.9f}")
    starts = 0.4 + jnp.arange(100_000) * 1e-6
    newton_all(starts).block_until_ready()
    roots = timed("newton, 100k starts", lambda: newton_all(starts).block_until_ready())
    print(f"{float(jnp.mean(roots)):.9f} {float(jnp.max(jnp.abs(jax.vmap(g)(roots)))):.0e}")


main()
gradients_torch.py (Python) 
"""The same derivatives under PyTorch: `torch.func.grad` and `vmap` over the starting points."""

import time

import torch
from torch.func import grad, vmap

torch.set_default_dtype(torch.float64)
# PyTorch's default is one thread per logical core; on a hybrid CPU that
# count spins on the efficiency cores and the batch takes hundreds of
# milliseconds some runs and seven others. The performance cores' count is
# the number to give it, as a program would.
torch.set_num_threads(min(8, torch.get_num_threads()))


def f(x, y):
    return torch.sin(x) * y + x * x


def g(x):
    return torch.exp(-x * x) * torch.cos(3.0 * x)


df = grad(f)
dg = grad(g)


def newton(x):
    for _ in range(6):
        x = x - g(x) / dg(x)
    return x


newton_all = vmap(newton)


def timed(label, run):
    best = 1e9
    answer = None
    for _ in range(5):
        started = time.perf_counter()
        answer = run()
        best = min(best, time.perf_counter() - started)
    print(f"# {label}: {best * 1000:.2f} ms")
    return answer


def main():
    x = torch.tensor(0.7)
    y = torch.tensor(2.0, requires_grad=True)
    v = f(x, y)
    grad_y = torch.autograd.grad(v, y)[0]
    print(f"{float(df(x, torch.tensor(2.0))):.9f} {float(v.detach() * 2.0 + grad_y):.9f}")
    starts = 0.4 + torch.arange(100_000, dtype=torch.float64) * 1e-6
    newton_all(starts)
    roots = timed("newton, 100k starts", lambda: newton_all(starts))
    print(f"{float(roots.mean()):.9f} {float(vmap(g)(roots).abs().max()):.0e}")


main()

Source: examples/36_autodiff.