Skip to content

← Flax step 83 of 100

Medium Primitives

nn.jit (Flax-aware JIT lift)

Jump to the task

Compile a Dense layer through Linen

Wrap nn.Dense with the lifted transformation nn.jit, then initialize and apply the resulting module. The layer still computes an affine map:

y = x @ kernel + bias

For input shape (D,) and output width F, the kernel has shape (D, F), the bias has shape (F,), and the result has shape (F,). There is no activation function or training step.

Transform the class, then instantiate it

The class-based form used here is:

JitDense = nn.jit(nn.Dense)
model = JitDense(features=F)
variables = model.init(rng, x)
out = model.apply(variables, x)

Pass the module class to the transform, not a constructed Dense instance. The lifted transformation handles Linen variables and random streams at the module-call boundary. Its transformed computation can be used during initialization as well as application.

For a pure call using existing variables, jax.jit(model.apply) is another valid way to compile the computation. Lifted nn.jit is useful when expressing compilation boundaries within a module hierarchy or composing Linen transformations. It is not required merely because parameters are present.

Compare the same weights

Compilation preserves the intended mathematical operation, but compiler optimizations can change floating-point rounding. Compare numerical results with a tolerance rather than assuming bit-for-bit identity.

This complete example supplies the same weights to three ways of evaluating the Dense layer:

import jax
import jax.numpy as jnp
import flax.linen as nn

x = jnp.array([1., 2.])
variables = {"params": {
    "kernel": jnp.array([[1., -1.], [2., .5]]),
    "bias": jnp.array([.5, 1.]),
}}
plain = nn.Dense(features=2)
lifted = nn.jit(nn.Dense)(features=2)
compiled_apply = jax.jit(plain.apply)
expected = jnp.array([5.5, 1.])

assert jnp.allclose(plain.apply(variables, x), expected)
assert jnp.allclose(lifted.apply(variables, x), expected)
assert jnp.allclose(compiled_apply(variables, x), expected)

The exercise initializes its weights from the supplied seed instead of using the explicit values above. Preserve the direct transformed module scope and initializer defaults to match its tests.

Compilation and caching

A JIT call may trace and compile before it executes. Later compatible calls can reuse a cached compilation. Different shapes, dtypes or static argument values can require another compilation; function identity and other compilation settings also matter. It is not simply “compile once forever.”

JIT is not an unconditional speed improvement. Compilation has a cost, and a small computation or a workload with frequent recompilation may not recover that cost. When measuring execution time, warm up the exact signature you will measure and wait for device work to finish, for example with result.block_until_ready().

Reuse transformed classes and compiled callables where practical in application code. Do not infer an exact compilation count just from where nn.jit(nn.Dense) appears in Python source. This exercise checks a seeded forward result, not cache hits or speed.

Problem

Implement nn_jit_forward(seed, x, features):

  1. Convert features to Python integer F and construct JitDense = nn.jit(nn.Dense).
  2. Instantiate it directly with features=F. Keep Dense's default kernel initializer, enabled bias and zero bias initializer.
  3. Initialize with jax.random.PRNGKey(int(seed)) and x.
  4. Apply those variables to the same x and return .reshape(-1).

x is a one-dimensional float array of length D; D and F are positive. Return a one-dimensional array of length F.

Root parameter keys are kernel and bias, containing D * F + F scalar parameters. Add no wrapper scope, explicit module name or extra layer. If JIT-compiling the whole entrypoint as well, keep seed and features static because the function converts them to Python integers.