Skip to content

← Flax NNX step 61 of 100

Medium Primitives

NNX nnx.jit

Jump to the task

Why this matters

A forward pass often repeats the same array operations on new inputs. nnx.jit lets JAX compile those operations together while handling the NNX model passed to the function. This can reduce Python dispatch overhead and enable compiler optimizations. Whether it helps depends on the workload; compilation itself also takes time.

Eager JAX already uses compiled operations. Adding JIT changes the size of the computation being compiled, rather than turning compilation on for the first time.

A compiled forward pass

import jax.numpy as jnp
from flax import nnx

model = nnx.Linear(3, 4, rngs=nnx.Rngs(0))
x = jnp.array([[1.0, 2.0, 3.0]], dtype=jnp.float32)

@nnx.jit
def forward(model, x):
    return model(x)

eager = model(x)
first = forward(model, x)
second = forward(model, x)
assert jnp.allclose(first, eager, rtol=1e-5, atol=1e-6)
assert jnp.allclose(second, first, rtol=1e-5, atol=1e-6)

model.kernel[...] = jnp.zeros_like(model.kernel[...])
updated = forward(model, x)
assert jnp.allclose(updated, model(x), rtol=1e-5, atol=1e-6)

Pass the model explicitly so each call receives its current parameters. The model does not have to be the first argument; that is just the convention used here. NNX also supports updates to Variable values inside nnx.jit, such as a counter increment, and propagates them back to the caller.

Compiled and eager calculations can differ slightly because floating-point operations may be optimized differently. Compare with a tolerance, rather than requiring bit-for-bit equality.

When can a compiled computation be reused?

Repeated calls to the same decorated function can reuse a compiled computation when the relevant structure and compilation settings match. Array shapes and dtypes matter, including those of the model's parameters. Static argument values and model structure can matter too. Changing only parameter values, while retaining their shapes and dtypes, normally does not require a new trace.

A new batch shape can require another specialization. Replacing a layer with an equivalent one does not necessarily change the model's structure. These are useful rules of thumb, not a complete formula for the cache key.

Define the decorated function once and reuse it. Creating a new function inside each iteration can prevent cache reuse. Python side effects such as print happen during tracing, so they are not a reliable way to log each execution; use jax.debug.print for runtime values.

Raw JAX transforms can also be used with NNX in supported patterns, including explicit graph/state handling. For this exercise, use nnx.jit with an explicit model argument so NNX handles the model's state.

Problem

Implement jit_forward(seed, x, features, num_calls):

  1. Build nnx.Linear(x.shape[-1], int(features), rngs=nnx.Rngs(int(seed))). Keep its default initializers and bias.
  2. Define a function fwd(model, x) decorated with @nnx.jit that returns model(x). Define it before the loop.
  3. Call it int(num_calls) times using the same model and input.
  4. Flatten and return the final output with out.reshape(-1).

This exercise creates the wrapper once per entrypoint call. In an application, keep a reusable forward function outside the repeated training or evaluation calls as well. The numerical tests check the final output; they cannot prove that you used JIT, repeated the calls, or obtained cache hits. Follow those implementation requirements even if an eager shortcut produces the same numbers.

Inputs:

  • seed: nonnegative integer-valued scalar, supplied as a float; cast to int.
  • x: finite float32 matrix of shape (N, D_in), with positive dimensions.
  • features: positive integer-valued scalar, supplied as a float; cast to int.
  • num_calls: positive integer-valued scalar, supplied as a float; cast to int. At least one call is needed to have a final output.

Output: a flat array of shape (N * int(features),).

The visualization illustrates input-shape specialization for a simple function. It does not describe every factor involved in NNX compilation.

Loading visualization…