Flax

An MLP regression trained with Flax (linen) and optax, converted from ordinary Python and checked under strict = true with nothing extra installed or configured. The jax plugin models the Flax and optax surface, so strict mode has a signature for every call.

Run it

python  train.ppy
ppy run train.ppy

What it prints

python train.ppy, ppy run train.ppy

device=cpu
loss 3.0810 -> 0.0105
learned

What the conversion typed

@ppy.pure
def target_curve(x: jax.Array) -> jax.Array:
    return 3.0 * x * x - 2.0 * x + 0.5

The untyped helpers got their types from their call sites. target_curve receives a jax.Array because that is what flows into it. The conversion also added @ppy.pure where the checker proved it.

What the conversion left alone

Two functions keep their author-written signatures:

• train_step, decorated with @partial(jax.jit, static_argnums=(0, 1))
• __call__, decorated with @nn.compact

A partial of a vouched decorator counts as that decorator, and the policy leaves a transformed function's annotations as they are.

A base class only the plugin knows

class Mlp(nn.Module) inherits init and apply from a base that is not in the project. The checker resolves them through the class's external MRO, which the plugin supplies, so model.init(key, xs) and model.apply(p, xs) type-check.

Layer constructors (nn.Dense), activations (nn.relu), optax.adam, and tx.update's (updates, state) pair all have signatures. Parameter pytrees are an explicit Any boundary declared in the source, rather than an inferred one.

Read on: Training with JAX · Plugins: JAX

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

29_flax/train.ppy

from functools import partial
from typing import Any

import jax
import jax.numpy as jnp
import optax
import ppy
from flax import linen as nn


class Mlp(nn.Module):
    hidden: int

    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        x = nn.Dense(self.hidden)(x)
        x = nn.relu(x)
        return nn.Dense(1)(x)


@ppy.pure
def target_curve(x: jax.Array) -> jax.Array:
    return 3.0 * x * x - 2.0 * x + 0.5


@ppy.pure
def relative_drop(first: float, last: float) -> float:
    return last / first


@partial(jax.jit, static_argnums=(0, 1))
def train_step(
    model: Mlp,
    tx: optax.GradientTransformation,
    params: Any,
    state: Any,
    xs: jax.Array,
    ys: jax.Array,
) -> Any:
    def loss_fn(p: Any) -> jax.Array:
        pred = model.apply(p, xs)
        return jnp.mean((pred - ys) ** 2)

    loss, grads = jax.value_and_grad(loss_fn)(params)
    updates, state = tx.update(grads, state)
    params = optax.apply_updates(params, updates)
    return params, state, loss


def main() -> None:
    key: jax.Array = jax.random.PRNGKey(0)
    xs: jax.Array = jnp.linspace(-1.0, 1.0, 256).reshape(-1, 1)
    ys: jax.Array = target_curve(xs)

    model: Mlp = Mlp(hidden=32)
    params: Any = model.init(key, xs)
    tx = optax.adam(1e-2)
    state: Any = tx.init(params)

    first: float = 0.0
    last: float = 0.0
    for i in range(200):
        params, state, loss = train_step(model, tx, params, state, xs, ys)
        if i == 0:
            first = float(loss)
        last = float(loss)
    print(f"device={jax.devices()[0].platform}")
    print(f"loss {first:.4f} -> {last:.4f}")
    print("learned" if relative_drop(first, last) < 0.05 else "did not learn")


main()

Source: examples/29_flax.