Tools for JAX

This repository implements a variety of tools for the differential programming library JAX.

Major components

Tjax’s major components are:

A dataclass decorator dataclass that facilitates defining structured JAX objects (so-called “pytrees”), which benefits from:
- the ability to mark fields as static (not available in chex.dataclass), and
- a display method that produces formatted text according to the tree structure.
A shim for the gradient transformation library optax that supports:
- easy differentiation and vectorization of “gradient transformation” (learning rule) parameters,
- gradient transformation objects that can be passed dynamically to jitted functions, and
- generic type annotations.
A pretty printer print_generic for aggregate and vector types, including dataclasses. (See display.) It features:
- support for traced values,
- colorized tree output for aggregate structures, and
- formatted tabular output for arrays (or statistics when there’s no room for tabular output).

Minor components

Tjax also includes:

Versions of custom_vjp and custom_jvp that support being used on methods: custom_vjp_method and custom_jvp_method (See shims.)
Tools for working with cotangents. (See cotangent_tools.)
JAX tree registration for NetworkX graph types. (See graph.)
A hashable immutable mapping frozendict that is registered as a JAX pytree, so it can be passed through transforms like jit and vmap. (See frozendict.)
Leaky integration leaky_integrate and Ornstein-Uhlenbeck process iteration diffused_leaky_integrate. (See leaky_integral.)
A testing function assert_tree_allclose that automatically produces testing code. And, a related function tree_allclose, which compares pytrees with matching structure. (See testing.)
Basic tools like divide_where. (See tools.)

Contribution guidelines

The implementation should be consistent with the surrounding style, be type annotated, and pass the linters below.

To run tests: uv run pytest

There are a few tools to clean and check the source:

uv run ruff check
uv run ruff format
uv run ty check

Metadata

Release files for tjax 1.8.0

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for tjax 1.8.0
File	Size	Uploaded
tjax-1.8.0.tar.gz	43.1 kB	Oct 7, 2026	Details

Built distribution (wheel)

Table of built distributions (wheels) for tjax 1.8.0
File	Interpreter	ABI	Platform	Reset
tjax-1.8.0-py3-none-any.whl 54.8 kB Oct 7, 2026	Python 3	none	any	Details

Total release size: 97.9 kB

Release files / tjax-1.8.0.tar.gz

Download URL	tjax-1.8.0.tar.gz
Size	43.1 kB
Tags	Source
SHA-256 checksum How to use checksums	`be1299167ce5bfb4c0ceb1bdd88f0c05f0ce05e95a8166013efb26988f2bd881`
BLAKE2b-256 checksum How to use checksums	`563401f038bf9a4f935f99e943d363984d728a8d57e047abd01447f0b14c5ba1`
Upload date	Oct 7, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	Yes
Uploaded via	`twine/7.0.0 CPython/3.13.14`

Provenance

Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.

Identity

Publishing platform: GitHub Actions
Publishing repository: github.com/NeilGirdhar/tjax
Publishing commit: github.com/NeilGirdhar/tjax/tree/9b9a5f8f411d0078d6e70324062e35372404db19

Workflow

Publishing configuration: github.com/NeilGirdhar/tjax/blob/9b9a5f8f411d0078d6e70324062e35372404db19/.github/workflows/release.yml
Publishing logs: github.com/NeilGirdhar/tjax/actions/runs/37586620206/attempts/1

Artifact

Subject: tjax-1.8.0.tar.gz
SHA-256 checksum: be1299167ce5bfb4c0ceb1bdd88f0c05f0ce05e95a8166013efb26988f2bd881 Verifying attestations

Release files / tjax-1.8.0-py3-none-any.whl

Download URL	tjax-1.8.0-py3-none-any.whl
Size	54.8 kB
Tags	Python 3
SHA-256 checksum How to use checksums	`fea31c2e15fcb9e96e29bab4a76dfdda6cd44801aab3a59bd8b350e510e874e2`
BLAKE2b-256 checksum How to use checksums	`ed21a33ee070481aec68c64dd7d6a102ed6f987e98edfc9f37f3f05bc0ee991c`
Upload date	Oct 7, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	Yes
Uploaded via	`twine/7.0.0 CPython/3.13.14`

Provenance

Identity

Publishing platform: GitHub Actions
Publishing repository: github.com/NeilGirdhar/tjax
Publishing commit: github.com/NeilGirdhar/tjax/tree/9b9a5f8f411d0078d6e70324062e35372404db19

Workflow

Publishing configuration: github.com/NeilGirdhar/tjax/blob/9b9a5f8f411d0078d6e70324062e35372404db19/.github/workflows/release.yml
Publishing logs: github.com/NeilGirdhar/tjax/actions/runs/37586620206/attempts/1

Artifact

Subject: tjax-1.8.0-py3-none-any.whl
SHA-256 checksum: fea31c2e15fcb9e96e29bab4a76dfdda6cd44801aab3a59bd8b350e510e874e2 Verifying attestations