tjax 0.3.5

Tools for JAX.

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

• A dataclass decorator that facilitates defining JAX trees, provides convenient text display, and provides a mypy plugin

• A custom VJP decorator that supports both static and non-differentiable arguments

• A random number generator class

• JAX tree registration for NetworkX graph types

• Testing tools that automatically produce testing code

See the documentation.

Contribution guidelines

• Conventions: PEP8.

• How to run tests: pytest .

• How to clean the source:

  • isort tjax

  • pylint tjax

  • mypy tjax

  • flake8 tjax