Training a JAX MLP

The same trainer as 21_training_torch with JAX, converted by ppy convert with no hand editing. Preprocessing goes from about 63 ms to under 1 ms. The training step stays where it was, about 25 ms for 100 steps either way, because XLA compiled it on its first call and there is no per-operator Python overhead left to remove.

Run it

python  train.ppy
ppy run train.ppy

What it prints

python train.ppy

# device: cpu
# native prep: False
prep      67.9 ms   checksum=-21433.891867
train     94.8 ms   loss 1.0663 -> 1.0109

ppy run train.ppy

# device: cpu
# native prep: True
prep       0.9 ms   checksum=-21433.891867
train    111.2 ms   loss 1.0663 -> 1.0109

Borrowed buffers, again

def standardize(raw: Buffer[float], out: Buffer[float], rows: int, cols: int) -> float:

The converter declared both parameters Buffer[float] because the body indexes them and never slices, and the values feeding them became array.array. The loop lowers natively.

@jax.jit is already the fast path

@jax.jit
def train_step(x, y, w1, b1, w2, b2, rate) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array]:
    grads = gradients_of(x, y, w1, b1, w2, b2)
    return (w1 - rate * grads[0], b1 - rate * grads[1], w2 - rate * grads[2], b2 - rate * grads[3])

ppy build can export a @jax.jit function whose inputs carry ppy.Shape and ppy.DType to StableHLO ahead of time, which saves the trace rather than the kernel (JAX export).

It declines to export a function that is differentiated, and says why: a serialized jax.export artifact carries no VJP. Routing forward_loss through one would break jax.grad on a program that runs correctly under plain CPython.

What the input had to get right

• standardize indexes rather than slices, as in the torch example.
• train_step takes and returns the parameters positionally rather than rebinding one list. Rebinding a variable from the return value of the function it is passed to makes its type self-referential, and inference gives up.
• forward_loss is called directly as well as through jax.grad. A function only ever reached through a higher-order transform has no call site to infer from.

Where the code comes from

Generated, not hand-written: train.ppy is exactly what ppy convert train.py --promote-buffers writes, and examples/verify_conversions.py checks that on every run.

Read on: Plugins: JAX · Flax

22_training_jax/train.ppy

import array
import math
import time

import jax
import jax.numpy as jnp
import ppy
from ppy import Buffer


def standardize(raw: Buffer[float], out: Buffer[float], rows: int, cols: int) -> float:
    total: float = 0.0
    for row in range(rows):
        base: int = row * cols
        target: int = row * cols * 2

        sum_: float = 0.0
        for i in range(cols):
            sum_ += raw[base + i]
        mean: float = sum_ / cols

        spread: float = 0.0
        for i in range(cols):
            spread += (raw[base + i] - mean) * (raw[base + i] - mean)
        deviation: float = math.sqrt(spread / cols) + 1e-8

        for i in range(cols):
            value: float = (raw[base + i] - mean) / deviation
            out[target + i] = value
            total += value
        for i in range(cols):
            interaction: float = out[target + i] * out[target + (i + 1) % cols]
            out[target + cols + i] = interaction
            total += interaction
    return total


@jax.jit
def forward_loss(
    x: jax.Array,
    y: jax.Array,
    w1: jax.Array,
    b1: jax.Array,
    w2: jax.Array,
    b2: jax.Array,
) -> jax.Array:
    hidden: jax.Array = jnp.maximum(jnp.dot(x, w1) + b1, 0.0)
    predicted: jax.Array = jnp.dot(hidden, w2) + b2
    residual: jax.Array = predicted - y
    return jnp.mean(residual * residual)


gradients_of = jax.grad(forward_loss, argnums=(2, 3, 4, 5))


@jax.jit
def train_step(
    x: jax.Array,
    y: jax.Array,
    w1: jax.Array,
    b1: jax.Array,
    w2: jax.Array,
    b2: jax.Array,
    rate: float,
) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array]:
    grads = gradients_of(x, y, w1, b1, w2, b2)
    return (
        w1 - rate * grads[0],
        b1 - rate * grads[1],
        w2 - rate * grads[2],
        b2 - rate * grads[3],
    )


def main() -> None:
    rows: int = 20000
    cols: int = 16

    key: jax.Array = jax.random.PRNGKey(0)
    k1, k2, k3, k4 = jax.random.split(key, 4)
    raw = array.array("d", jax.random.normal(k1, (rows * cols,)).tolist())
    features = array.array("d", [0.0] * (rows * cols * 2))

    started: float = time.perf_counter()
    checksum: float = standardize(raw, features, rows, cols)
    prep_ms: float = (time.perf_counter() - started) * 1000.0

    x: jax.Array = jnp.asarray(features, dtype=jnp.float32).reshape(rows, cols * 2)
    y: jax.Array = jax.random.normal(k2, (rows, 1))
    w1 = jax.random.normal(k3, (cols * 2, 32)) * 0.1
    b1 = jnp.zeros((32,))
    w2 = jax.random.normal(k4, (32, 1)) * 0.1
    b2 = jnp.zeros((1,))

    first: float = float(forward_loss(x, y, w1, b1, w2, b2))
    w1, b1, w2, b2 = train_step(x, y, w1, b1, w2, b2, 0.02)
    w1.block_until_ready()

    started = time.perf_counter()
    for _ in range(100):
        w1, b1, w2, b2 = train_step(x, y, w1, b1, w2, b2, 0.02)
    last: float = float(forward_loss(x, y, w1, b1, w2, b2))
    train_ms: float = (time.perf_counter() - started) * 1000.0

    print(f"# device: {jax.devices()[0].platform}")
    print(f"# native prep: {ppy.native.compiled(standardize)}")
    print(f"prep  {prep_ms:8.1f} ms   checksum={checksum:.6f}")
    print(f"train {train_ms:8.1f} ms   loss {first:.4f} -> {last:.4f}")


main()

Source: examples/22_training_jax.