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.1

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.1
File	Size	Uploaded
tjax-1.8.1.tar.gz	43.1 kB	Oct 7, 2026	Details

Built distribution (wheel)

Table of built distributions (wheels) for tjax 1.8.1
File	Interpreter	ABI	Platform	Reset
tjax-1.8.1-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.1.tar.gz

Download URL	tjax-1.8.1.tar.gz
Size	43.1 kB
Tags	Source
SHA-256 checksum How to use checksums	`f82ca6fdba5f5ad78f693b9f943b528a40cc01461d535a9af90127fe33aa9406`
BLAKE2b-256 checksum How to use checksums	`48c87607e17d31900254d99301a110c425589dc820fcf4fa204d30c84dd21d82`
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/bba10387f94c5dc0a80e6f2114d598f0b1d80c17

Workflow

Publishing configuration: github.com/NeilGirdhar/tjax/blob/bba10387f94c5dc0a80e6f2114d598f0b1d80c17/.github/workflows/release.yml
Publishing logs: github.com/NeilGirdhar/tjax/actions/runs/37593451906/attempts/1

Artifact

Subject: tjax-1.8.1.tar.gz
SHA-256 checksum: f82ca6fdba5f5ad78f693b9f943b528a40cc01461d535a9af90127fe33aa9406 Verifying attestations

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

Download URL	tjax-1.8.1-py3-none-any.whl
Size	54.8 kB
Tags	Python 3
SHA-256 checksum How to use checksums	`3be1e8ead94a8869646e4a29e930935ec33bc1b99cd27c070a5e8ffdb4d79c30`
BLAKE2b-256 checksum How to use checksums	`a35a1573a762307fb6d3a1e1c7209242e9cf3e6e802584d3f70ff55b2b23337d`
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/bba10387f94c5dc0a80e6f2114d598f0b1d80c17

Workflow

Publishing configuration: github.com/NeilGirdhar/tjax/blob/bba10387f94c5dc0a80e6f2114d598f0b1d80c17/.github/workflows/release.yml
Publishing logs: github.com/NeilGirdhar/tjax/actions/runs/37593451906/attempts/1

Artifact

Subject: tjax-1.8.1-py3-none-any.whl
SHA-256 checksum: 3be1e8ead94a8869646e4a29e930935ec33bc1b99cd27c070a5e8ffdb4d79c30 Verifying attestations