jax-dataclasses

What it is and what it does

jax_dataclasses is a thin wrapper around Python's standard dataclasses that integrates them seamlessly into JAX's pytree system. It automatically registers decorated classes as pytree nodes, making them usable at JAX API boundaries (e.g., as function arguments to jitted code). The package adds support for marking fields as static—meaning they are constant at compile time and won't be traced—and enables serialization through flax.serialization.

The main value proposition is ergonomic: unlike hand-registering pytrees, you write normal dataclass syntax and get pytree behavior automatically. The package also provides copy_and_mutate(), a context manager that temporarily unfreezes nested dataclass structures for easier in-place modifications, addressing a common pain point when working with deeply nested immutable objects in JAX.

Use it for:

• Define model parameters and state as frozen dataclasses in JAX neural network code, automatically compatible with jit and vmap.
• Mark configuration fields as static to prevent JAX from recompiling when only non-array metadata changes.
• Serialize and deserialize JAX model checkpoints via flax.serialization without manual pytree registration.
• Modify deeply nested dataclass structures in a readable way using copy_and_mutate instead of chained dataclasses.replace() calls.
• Use standard type checkers and IDE autocomplete on JAX dataclasses since the decorator matches the standard library API.

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

Provides a JAX-compatible wrapper around Python dataclasses that automatically registers them as pytrees and enables serialization, while supporting static fields and nested mutations.

Yes, if you are building JAX applications with structured state or model parameters. The package solves a real integration gap between Python dataclasses and JAX's pytree system with minimal overhead. Install friction is low and the MIT license is unrestrictive. The aging maintenance status (238 days since last release) is a minor concern but not a blocker—the package is stable, unarchived, and has no known vulnerabilities.

Install

jax-dataclasses on PyPI

pip install jax-dataclasses

uv add jax-dataclasses

poetry add jax-dataclasses

Installing jax-dataclasses

Before you install

Low friction: pure Python wheel with three runtime dependencies (jax, jaxlib, typing_extensions). Maintenance status is aging—last release was 238 days ago—but the repository remains active and unarchived.

License in practice

MIT license (permissive): you can use, modify, and distribute this package freely in commercial and private projects with minimal restrictions.

Quickstart

pip install jax_dataclasses

import jax_dataclasses as jdc
import jax

@jdc.pytree_dataclass
class Model:
    params: jax.Array
    name: jdc.Static[str]

model = Model(params=jax.numpy.zeros(10), name="test")

Requires Python >=3.9 and jax/jaxlib installed.

Verify before relying

• Whether flax.serialization integration works with all flax versions or requires a specific version.
• Performance characteristics when working with very deeply nested dataclass structures.
• Compatibility with JAX's latest pytree API changes beyond Python 3.12.

Package facts

Evidence: jax_dataclasses-1.6.3-py3-none-any.whl

Tags

More Artificial Intelligence packages