Skip to main content

Tools for JAX.

Project description

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

Major components

Tjax’s major components are:

  • A dataclass decorator dataclasss that facilitates defining JAX trees, and has a MyPy plugin. (See dataclass and mypy_plugin.)

  • A fixed point finding library heavily based on fax. Our library supports stochastic iterated functions, and avoids leaking JAX tracers. (See fixed_point.)

Minor components

Tjax also includes:

  • An object-oriented wrapper on top of optax. (See gradient.)

  • A pretty printer print_generic for aggregate and vector types, including dataclasses. (See display.)

  • Versions of custom_vjp and custom_jvp that support both static and non-differentiable arguments. (See shims.)

  • Tools for working with cotangents: copy_cotangent and print_cotangent. (See cotangent_tools.)

  • A random number generator class Generator. (See generator.)

  • JAX tree registration for NetworkX graph types. (See graph.)

  • Leaky integration leaky_integrate and Ornstein-Uhlenbeck process iteration diffused_leaky_integrate. (See leaky_integral.)

  • An improved version of jax.tree_util.Partial. (See partial.)

  • A Matplotlib trajectory plotter PlottableTrajectory. (See plottable_trajectory.)

  • A testing function assert_jax_allclose that automatically produces testing code, and the related function jax_allclose. (See testing.)

  • Basic tools sum_tensors and is_scalar. (See tools.)

Also, 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.6

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

Uploaded Source

Built Distribution

tjax-0.6-py3-none-any.whl (36.6 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