Registry / ai-ml / jax-dataclasses

jax-dataclasses

library1.6.3pypypi✓ verified 83d ago

A library that provides a dataclass-like decorator for use with JAX, enabling mutable-style syntax with functional transformations, static fields, and support for pytree nodes. Current version is 1.6.3, actively maintained, with releases every few months.

ai-ml serialization

pip install jax-dataclasses

github ↗package ↗

INSTALL

IMPORT

SIG · JAX-DATACLASSES

jax-dataclasses

ai-mlpythonv1.6.3

Install

11.7s avg

Import

1964ms

Disk

606MB

Pass rate

5/ 10

Env Coverage5 / 10

glibc

3.9–3.13

musl

3.9–3.13

Install & Compatibility

Where this runs

tested against v1.6.3 · pip install

no network on importno background threads

Install × environment matrix

Each cell = how many times install + import succeeded across repeated harness runs. Partial = flaky.
glibc = Debian/Ubuntu slim · musl = Alpine Linux

musl

py 3.10–3.95 runs

build_error

glibc

py 3.10–3.95 runs

installs and imports cleanly · install 11.7s · import 1.964s · 585MB

606MB installed

● package 606MB

Code

Verified usage

Verified import paths — ran on the pinned version, not inferred.

jdc

✓ import jax_dataclasses as jdc

✗ from jax_dataclasses import jdc

jdc is the module alias, not a submodule. Using 'from ... import jdc' will raise ImportError.

jdc.jit

✓ import jax_dataclasses as jdc; @jdc.jit

✗ from jax_dataclasses import jit

jit is an attribute of the module, not importable directly. Use jdc.jit.

Static

✓ from jax_dataclasses import Static

Creates a simple pytree dataclass and demonstrates functional mutation via jdc.replace.

import jax import jax_dataclasses as jdc @jdc.pytree_dataclass class MyModel: a: jax.Array b: jax.Array model = MyModel(a=jax.numpy.array(1.0), b=jax.numpy.array(2.0)) # Functional update new_model = jdc.replace(model, a=jax.numpy.array(3.0)) print(new_model.a, new_model.b) # 3.0, 2.0

Debug

Known issues

deprecatedThe shape / datatype annotation API (e.g., @jdc.pytree_dataclass(shape_dtype=...)) is deprecated since v1.6.0. Use Static[] annotations instead.

fix

Replace shape_dtype annotations with jdc.Static[] for static fields.

affects: >=1.6.0

gotchaDo not use standard Python dataclass decorator (from dataclasses import dataclass) on a class with JAX arrays; it will break pytree registration. Always use @jdc.pytree_dataclass.

fix

Use @jdc.pytree_dataclass or @jdc.pytree_dataclass(frozen=True).

affects: all

breakingIn v1.6.2, Python 3.8 support was dropped. Requires Python >=3.9.

fix

Ensure Python >= 3.9.

affects: >=1.6.2

gotchaStatic field annotations must use jdc.Static[] (e.g., a: jdc.Static[int]) to be properly handled; using typing.ClassVar may not work correctly.

fix

Use jdc.Static[type] for static fields.

affects: >=1.6.0

Errors

Common errors & fixes

AttributeError: module 'jax_dataclasses' has no attribute 'jdc'

Attempting to import jdc as a submodule: from jax_dataclasses import jdc

fix

Use import jax_dataclasses as jdc

TypeError: replace() got an unexpected keyword argument 'a'

Using jdc.replace on a class not decorated with @jdc.pytree_dataclass (maybe used standard dataclass).

fix

Ensure the class is decorated with @jdc.pytree_dataclass.

Upgrade

Version history

1.6.3latest on PyPI · released Dec 19, 2025

Audit

Dependencies

jaxrequiredRequired for pytree registration and JIT compilation.

chexoptionalUsed for testing; optional for runtime.

Agent activity

12 hits · last 30 days

node

Resources

githubhttp://github.com/brentyi/jax_dataclasses ↗

packagehttps://pypi.org/project/jax-dataclasses/ ↗