Package: jax-dataclasses
jax_dataclassesOverviewInstallationCore interfaceStatic fieldsMutationsAlternativesMiscOverviewjax_dataclassesprovides a simple wrapper arounddataclasses.dataclassfor use in
JAX, which enables automatic support for:Pytreeregistration. This
allows dataclasses to be used at API boundaries in JAX.Serialization viaflax.serialization.Distinguishing features include:An annotation-based interface for marking static fields.Improved ergonomics for "model surgery" in nested structures.InstallationIn Python >=3.7:pipinstalljax_dataclassesWe can then import:importjax_dataclassesasjdcCore interfacejax_dataclassesis meant to provide a drop-in replacement fordataclasses.dataclass:jdc.pytree_dataclasshas
the same interface asdataclasses.dataclass, but also registers the target
class as a pytree node.We also provide several aliases:jdc.[field, asdict, astuples, is_dataclass, replace]are identical to
their counterparts in the standard dataclasses library.Static fieldsTo mark a field as static (in this context: constant at compile-time), we can
wrap its type withjdc.Static[]:@jdc.pytree_dataclassclassA:a:jax.Arrayb:jdc.Static[bool]In a pytree node, static fields will be treated as part of the treedef instead
of as a child of the node; all fields that are not explicitly marked static
should contain arrays or child nodes.Bonus: if you likejdc.Static[], we also introducejdc.jit(). This enables use in function
signatures, for example:@jdc.jitdeff(a:jax.Array,b:jdc.Static[bool])->jax.Array:...MutationsAll dataclasses are automatically marked as frozen and thus immutable (even when
nofrozen=parameter is passed in). To make changes to nested structures
easier,jdc.copy_and_mutate(a) makes a copy of a
pytree and (b) returns a context in which any of that copy's contained
dataclasses are temporarily mutable:importjaxfromjaximportnumpyasjnpimportjax_dataclassesasjdc@jdc.pytree_dataclassclassNode:child:jax.Arrayobj=Node(child=jnp.zeros(3))withjdc.copy_and_mutate(obj)asobj_updated:# Make mutations to the dataclass. This is primarily useful for nested# dataclasses.## Does input validation by default: if the treedef, leaf shapes, or dtypes# of `obj` and `obj_updated` don't match, an AssertionError will be raised.# This can be disabled with a `validate=False` argument.obj_updated.child=jnp.ones(3)print(obj)print(obj_updated)AlternativesA few other solutions exist for automatically integrating dataclass-style
objects into pytree structures. Great ones include:chex.dataclass,flax.struct, andtjax.dataclass. These all influenced
this library.The main differentiators ofjax_dataclassesare:Static analysis support.tjaxhas a custom mypy plugin to enable type
checking, but isn't supported by other tools.flax.structimplements thedataclass_transformspec proposed by pyright, but isn't supported by other tools. Because@jdc.pytree_dataclasshas the same API as@dataclasses.dataclass, it can
include pytree registration behavior at runtime while being treated as the
standard decorator during static analysis. This means that all static
checkers, language servers, and autocomplete engines that support the standarddataclasseslibrary should work out of the box withjax_dataclasses.Nested dataclasses.Making replacements/modifications in deeply nested
dataclasses can be really frustrating. The three alternatives all introduce a.replace(self, ...)method to dataclasses that's a bit more convenient than
the traditionaldataclasses.replace(obj, ...)API for shallow changes, but
still becomes really cumbersome to use when dataclasses are nested.jdc.copy_and_mutate()is introduced to address this.Static field support.Parameters that should not be traced in JAX should
be marked as static. This is supported inflax,tjax, andjax_dataclasses, but notchex.Serialization.When working withflax, being able to serialize
dataclasses is really handy. This is supported inflax.struct(naturally)
andjax_dataclasses, but notchexortjax.You can also eschew the dataclass-style interface entirely;see how brax registers pytrees.
This is a reasonable thing to prefer: it requires some floating strings and
breaks things that I care about but you may not (like immutability and__post_init__), but gives more flexibility with custom__init__methods.Miscjax_dataclasseswas originally written for and factored out ofjaxfg, whereNick Heppertprovided valuable feedback.