Skip to content

← JAX Conceptual Deep-Dives step 1 of 15

Medium Primitives

register_pytree_node

Jump to the task

Why this matters

A plain Python class is normally one opaque pytree leaf, even when its attributes contain arrays. Registering it tells JAX how to reach those attributes and rebuild the object.

A NamedTuple already has pytree support: its fields are children by default. Ordinary dataclasses need registration, for example with jax.tree_util.register_dataclass. The example below uses a plain class so you can see what manual registration adds.

Worked mini-example

import jax
import jax.numpy as jnp

class Pair:
    def __init__(self, first, second):
        self.first = first
        self.second = second

jax.tree_util.register_pytree_node(
    Pair,
    lambda pair: ((pair.first, pair.second), None),
    lambda aux, children: Pair(*children),
)

pair = Pair(jnp.array([1.0, 2.0]), jnp.array([3.0, 4.0]))
leaves = jax.tree_util.tree_leaves(pair)  # [pair.first, pair.second]

@jax.jit
def total(pair):
    return jnp.sum(pair.first) + jnp.sum(pair.second)

result = total(pair)  # 10

Flattening exposes the two children. Unflattening receives the auxiliary data first and the children second, then reconstructs a Pair. Here there is no metadata, so auxiliary data is None.

Common pitfalls

  • Children can themselves be nested pytrees. They are not limited to individual arrays.
  • Keep reconstruction simple. Transformations may supply tracers or placeholder objects, so avoid converting or validating array values in the constructor used by unflattening.
  • Register a class once, normally when its module loads, rather than on each call to a numerical function.
  • Metadata belongs in aux_data and needs stable equality and hashing. Array data belongs in children.
  • Registration exposes structure. The function must still obey each transformation's rules for shapes, dtypes and operations.

Problem

Implement sum_packed_named_pair(x_first, x_second). The two one-dimensional JAX arrays arrive as separate arguments. Return the sum of all their elements as a scalar.

The arrays may have different lengths. For [1, 2] and [3, 4, 5], return 15. An empty array contributes zero.

Sum each array separately, then add the scalar sums. Adding the arrays element by element is not the requested operation.

This exercise checks the numerical operation, not registration. The worked example shows how the same operation can accept a registered container. The diagram illustrates general nesting.

Loading visualization…