JAX Dataclass Utility

JaxDataclass: a frozen dataclass mixin that registers as a JAX pytree.

Motivation: flax.struct.dataclass registers pytrees using __init__ in tree_unflatten. When beartype/jaxtyping are active (via install_import_hook), they instrument __init__ with runtime type checks. During JAX pytree reconstruction inside jit/vmap/scan, JAX may pass None sentinels or abstract tracer values, which fail the type checks.

This mixin uses object.__new__ in tree_unflatten, bypassing __init__ entirely during reconstruction, so beartype never sees the intermediate values.

See: https://github.com/patrick-kidger/jaxtyping/issues/177

class nnx_ppo.jax_dataclass.JaxDataclass

Bases: object

Frozen dataclass mixin that registers as a JAX pytree.

Usage:

@jax.tree_util.register_pytree_node_class @dataclasses.dataclass(frozen=True) class MyStruct(JaxDataclass):

x: Float[Array, “batch”] y: int

replace(**kwargs)

tree_flatten()

classmethod tree_unflatten(names, children)