tjax · PyPI

Tools for JAX

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

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

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

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:

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

tjax-1.7.1.tar.gz (42.8 kB view details)

Uploaded Jul 14, 2026 Source

If you're not sure about the file name format, learn more about wheel file names.

The dropdown lists show the available interpreters, ABIs, and platforms. Enable javascript to be able to filter the list of wheel files.

tjax-1.7.1-py3-none-any.whl (54.6 kB view details)

Uploaded Jul 14, 2026 Python 3

Details for the file tjax-1.7.1.tar.gz.

The following attestation bundles were made for tjax-1.7.1.tar.gz:

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Details for the file tjax-1.7.1-py3-none-any.whl.

Hashes for tjax-1.7.1-py3-none-any.whl
Algorithm	Hash digest
SHA256	`d9882b5b7a6116162fb863ff5a97b22d6a6622633fd8fa2382a1f4f33c08d9b1`
MD5	`2a1a7d6591435c71c876a1201e487ed1`
BLAKE2b-256	`a360d5e8f5bcc33a134365f451d3db31ccb2c32b26fdc32686fde98abeab4722`

The following attestation bundles were made for tjax-1.7.1-py3-none-any.whl:

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.