Skip to content

← JAX Conceptual Deep-Dives step 3 of 15

Hard Primitives

Pytree Aux Data

Jump to the task

Why this matters

A custom pytree can keep array data separate from static configuration. Its flattening function returns (children, aux_data): children are traversed as data, while auxiliary data stays in the tree's structure.

For example, scaling a layer's parameter arrays should not also try to multiply its text name. Storing the name as auxiliary data keeps it out of the leaf list.

Worked mini-example

import jax
import jax.numpy as jnp

class LayerParams:
    def __init__(self, weight, bias, name):
        self.weight = weight
        self.bias = bias
        self.name = name

jax.tree_util.register_pytree_node(
    LayerParams,
    lambda layer: ((layer.weight, layer.bias), layer.name),
    lambda name, children: LayerParams(*children, name),
)

layer = LayerParams(jnp.array([1.0, 2.0]), jnp.array([0.5]), "linear")
scaled = jax.tree_util.tree_map(lambda leaf: leaf * 2, layer)
# scaled.weight = [2, 4], scaled.bias = [1], scaled.name = "linear"

The name survives reconstruction but is not passed to the scaling function. Only the two arrays appear in tree_leaves(layer).

Common pitfalls

  • Auxiliary data must support meaningful equality and hashing. Use immutable metadata such as strings or tuples of simple values. Keep arrays in the children.
  • Changing auxiliary data changes the tree definition and can require a new compiled version. Returning to earlier metadata may reuse a cached version.
  • A non-trainable number need not be static. A scale or learning rate that changes frequently can be a dynamic scalar input or child, avoiding specialization on every value.
  • Gradient transformations operate on eligible child leaves; metadata is carried along rather than differentiated.
  • Use None when there is no auxiliary data. This is distinct from putting a field named "aux" in an ordinary dictionary.

Problem

Implement aux_data_demo(x, scale) to return x * scale. The input x is a one-dimensional floating-point JAX array; scale is a scalar. Return an array with the same shape as x.

For x = [1, 2] and scale = 3, return [3, 6]. Zero scale gives zeros for finite inputs; empty input stays empty.

These arguments are passed directly. There is no custom node in this exercise, and scale is not automatically auxiliary data. The tests check multiplication; the worked example demonstrates actual metadata handling. The diagram shows ordinary tree nodes, not a custom registration.

Loading visualization…