Skip to content

← JAX Performance & Distributed step 4 of 25

Easy Primitives

jit with Pytree Input

Jump to the task

Why this matters

A jitted function can accept a dictionary of model parameters and an input array together. You do not need to concatenate the parameters into one vector first: JAX understands the dictionary as a pytree.

JAX tracks the container structure and the types of its dynamic leaves. It traces array operations on those leaves and can return another nested structure.

Worked mini-example

import jax
import jax.numpy as jnp

@jax.jit
def forward(params, x):
    return params["w"] @ x + params["b"]

params = {
    "w": jnp.array([[1.0, 2.0], [3.0, 4.0]]),
    "b": jnp.array([0.5, -0.5]),
}
y = forward(params, jnp.array([1.0, 2.0]))  # [5.5, 10.5]

New weights with the same shapes and dtypes can reuse the compiled version. Changing the structure, a leaf's shape or dtype, or static metadata can require another version.

Common pitfalls

  • Reordering insertion of the same keys in an ordinary dictionary does not change its pytree structure. JAX traverses sorted keys.
  • Python numeric scalars are normally dynamic JIT inputs, including inside pytrees. Changing a float's value does not automatically trigger retracing.
  • A string leaf is not automatically static. JIT generally rejects it as a dynamic input. Put configuration strings in registered static metadata, or use an explicit static argument where appropriate.
  • Custom classes need suitable pytree support. Registration also does not make unsupported dynamic leaf types valid JIT inputs.

Problem

Implement jit_pytree_sum(packed_inputs). Define a reusable helper with jax.jit(jnp.sum) outside the entrypoint, then call it on the supplied one-dimensional JAX array.

Return its element sum as a scalar. For [2, 3, 5], return 10. An empty input returns zero.

A single array is a pytree with one leaf. This exercise uses that simplest case; its numeric tests do not check nested traversal or prove that JIT was used. The worked example shows a dictionary input, and the diagram shows a separate nested structure.

Loading visualization…