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.9.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.9.0
File	Size	Uploaded
tjax-1.9.0.tar.gz	45.6 kB	Oct 10, 2026	Details

Built distribution (wheel)

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

Total release size: 100.5 kB

Release files / tjax-1.9.0.tar.gz

Download URL	tjax-1.9.0.tar.gz
Size	45.6 kB
Tags	Source
SHA-256 checksum How to use checksums	`63fb4021ffc4efcd93f3f6c1d2a211e197ba3119f30469fd9b6e344de83c7cf1`
BLAKE2b-256 checksum How to use checksums	`4d929b4d7877803e5b88478a8eb0bdd7460f658a03187c0742640d027837a4a9`
Upload date	Oct 10, 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/2209af8188afc549f0eb1d48c4c1c8b28e0bfc16

Workflow

Publishing configuration: github.com/NeilGirdhar/tjax/blob/2209af8188afc549f0eb1d48c4c1c8b28e0bfc16/.github/workflows/release.yml
Publishing logs: github.com/NeilGirdhar/tjax/actions/runs/38013944999/attempts/1

Artifact

Subject: tjax-1.9.0.tar.gz
SHA-256 checksum: 63fb4021ffc4efcd93f3f6c1d2a211e197ba3119f30469fd9b6e344de83c7cf1 Verifying attestations

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

Download URL	tjax-1.9.0-py3-none-any.whl
Size	54.9 kB
Tags	Python 3
SHA-256 checksum How to use checksums	`252ec066e117a85618abee5bc56890c63ad71959c58f555d30e8a678386ed594`
BLAKE2b-256 checksum How to use checksums	`9570296a4f709be4e6bd809d7e0720b8528004afdf88a572db3220481c4d5b77`
Upload date	Oct 10, 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/2209af8188afc549f0eb1d48c4c1c8b28e0bfc16

Workflow

Publishing configuration: github.com/NeilGirdhar/tjax/blob/2209af8188afc549f0eb1d48c4c1c8b28e0bfc16/.github/workflows/release.yml
Publishing logs: github.com/NeilGirdhar/tjax/actions/runs/38013944999/attempts/1

Artifact

Subject: tjax-1.9.0-py3-none-any.whl
SHA-256 checksum: 252ec066e117a85618abee5bc56890c63ad71959c58f555d30e8a678386ed594 Verifying attestations