jax-dataclasses
Dataclasses + JAX
Decision gist · record as of 2026-08-14
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.AI-flagged interpretation of the facts on this page — verify before relying
Before you install
- Requires Python >=3.9 and jax/jaxlib installed.
- 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 · maintenance · safety
MIT (permissive) — MIT license (permissive): you can use, modify, and distribute this package freely in commercial and private projects with minimal restrictions.
last release 2025-12-19 (238 days) · last repo commit 2025-12-19 · 76 stars
0 known vulnerabilities (OSV.dev, 2026-08-14) · 103,209 downloads/mo, #12,821 on PyPI
Alternatives
Verify before relying
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")- 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.
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 on it.
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
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.
Requires Python >=3.9 and jax/jaxlib installed.
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")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 the current Python release >=3.9 |
| Install friction | Low. Pure-Python wheel |
| Runtime dependencies | 3 packagesjaxjaxlibtyping_extensions |
| Maintenance | Aging 238 days since the last release |
| Last repo commit | |
| First released | |
| Downloads | 103,209 / month, #12,821 on PyPI 30-day window, as of 2026-08-14 |
| Known vulnerabilities | None known OSV.dev, checked 2026-08-14 |
| Classifiers | License :: OSI Approved :: MIT LicenseOperating System :: OS IndependentProgramming Language :: Python :: 3.10Programming Language :: Python :: 3.11Programming Language :: Python :: 3.12Programming Language :: Python :: 3.9 |
Evidence: jax_dataclasses-1.6.3-py3-none-any.whl
Tags
Let your AI agent find packages like this
Example. Real query, live index.
You found this page by searching. An agent finds it by wishing: SkillFed indexes 14,416 PyPI packages by what they can do, searchable in plain language.
wish › “jax pytree dataclass”
- jax-dataclassesProvides a JAX-compatible wrapper around Python dataclasses that…
- equinoxEquinox provides neural network and model building on top of JAX with…
- jaxliejaxlie implements Lie groups (SO2, SE2, SO3, SE3) for rigid body…
Give your agent the search over MCP, or paste the wish link into any chat.
More Artificial Intelligence packages
LiteLLM provides a unified Python interface to call 100+ LLM providers (OpenAI, Anthropic, Gemini, Bedrock, Azure, and others) using OpenAI-compatible API format, available as both a Python SDK and a self-hosted AI Gateway proxy server.
Install it if you need to work with multiple LLM providers or want to centralize LLM routing in your organization.
Client library and CLI tool for downloading, uploading, and managing models, datasets, and repositories on the Hugging Face Hub platform.
Install it if you work with Hugging Face Hub models or datasets.
LangChain provides a framework for building agents and LLM-powered applications by composing language models, tools, and memory through a unified API that abstracts over multiple model providers.
hf-xet provides chunk-based deduplication and efficient file transfer for the Hugging Face Hub, enabling faster uploads and downloads of large files with local disk caching.
Tokenizers converts raw text into token sequences for NLP models, with support for training custom vocabularies and using pre-built tokenizers (BPE, WordPiece) optimized for speed via Rust.
Transformers provides a unified framework for loading, fine-tuning, and running state-of-the-art pretrained models across text, vision, audio, video, and multimodal tasks using PyTorch, JAX, or TensorFlow.
Install it if you need to run or train any transformer-based model for NLP, vision, audio, or multimodal tasks.
See also dataclass-wizard · databind.json · djangorestframework-dataclasses · equinox · databind · itemadapter · serpyco-rs · dataclasses-json · databind.core · typed-json-dataclass