Tools for JAX.
Project description
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
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
tjax-0.3.4.tar.gz
(17.8 kB
view hashes)
Built Distribution
tjax-0.3.4-py3-none-any.whl
(24.2 kB
view hashes)