---
id: jax-dataclasses
version: "1.6.3"
license: MIT
license_treatment: permissive
maintenance: aging
---

# jax-dataclasses — Dataclasses + JAX

License: permissive · Maintenance: aging · Downloads: 103.2K/mo

## 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 above — 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

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

- License: MIT (permissive)
- Python support: supports_current
- Install friction: low
- Maintenance: aging
- Downloads: 103.2K/month (top 15,000 on PyPI)
- Known vulnerabilities: none known

## Tags

jax pytree dataclass, jax dataclass registration, frozen dataclass jax, static fields jax, jax nested dataclass mutations, jax serializable dataclass, jax pytree wrapper, jax-integration, pytree-registration, immutable-dataclass

[View on SkillFed](https://skillfed.io/packages/jax-dataclasses) · [View on PyPI](https://pypi.org/project/jax-dataclasses/)