jax.tree_util.register_pytree_node_class

jax.tree_util.register_pytree_node_class(cls)

Extends the set of types that are considered internal nodes in pytrees.

This function is a thin wrapper around register_pytree_node, and provides a class-oriented interface.

Parameters:

cls (Typ) – a type to register as a pytree. The class must define a tree_flatten method and a tree_unflatten class method, which implement flatten_func and unflatten_func as described in jax.tree_util.register_pytree_node().

Returns:

The input class cls is returned unchanged after being added to JAX’s pytree registry. This return value allows register_pytree_node_class to be used as a decorator.

Return type:

Typ

See also

• register_static(): simpler API for registering a static pytree.

• register_dataclass(): simpler API for registering a dataclass.

• register_pytree_node()

• register_pytree_with_keys()

• register_pytree_with_keys_class()

Examples

Here we’ll define a custom container that will be compatible with jax.jit() and other JAX transformations:

>>> import jax
>>> @jax.tree_util.register_pytree_node_class
... class MyContainer:
...   def __init__(self, name, x, y):
...     self.name = name  # static metadata
...     self.x = x        # array data
...     self.y = y        # array data
...   def tree_flatten(self):
...     children = (self.x, self.y)
...     aux_data = self.name
...     return (children, aux_data)
...   @classmethod
...   def tree_unflatten(cls, aux_data, children):
...     return cls(aux_data, *children)
...
>>> m = MyContainer('m', jnp.zeros(4), jnp.arange(4))
>>> def f(m):
...   return m.x + 2 * m.y
>>> jax.jit(f)(m)
Array([0., 2., 4., 6.], dtype=float32)