Skip to main content

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

This version

0.5.3

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.5.3.tar.gz (27.0 kB view hashes)

Uploaded Source

Built Distribution

tjax-0.5.3-py3-none-any.whl (35.4 kB view hashes)

Uploaded Python 3

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page