Skip to main content

JAX + Units

Project description

JAX + Units

Built with JAX and Pint!

This module provides an interface between JAX and Pint to allow JAX to support operations with units. The propagation of units happens at trace time, so jitted functions should see no runtime cost. This library is experimental so expect some sharp edges.

For example:

>>> import jax
>>> import jax.numpy as jnp
>>> import jpu
>>>
>>> u = jpu.UnitRegistry()
>>>
>>> @jax.jit
... def add_two_lengths(a, b):
...     return a + b
...
>>> add_two_lengths(3 * u.m, jnp.array([4.5, 1.2, 3.9]) * u.cm)
<Quantity([3.045 3.012 3.039], 'meter')>

Installation

To install, use pip:

python -m pip install jpu

The only dependencies are jax and pint, and these will also be installed, if not already in your environment. Take a look at the JAX docs for more information about installing JAX on different systems.

Usage

Here is a slightly more complete example:

>>> import jax
>>> import numpy as np
>>> from jpu import UnitRegistry, numpy as jnpu
>>>
>>> u = UnitRegistry()
>>>
>>> @jax.jit
... def projectile_motion(v_init, theta, time, g=u.standard_gravity):
...     """Compute the motion of a projectile with support for units"""
...     x = v_init * time * jnpu.cos(theta)
...     y = v_init * time * jnpu.sin(theta) - 0.5 * g * jnpu.square(time)
...     return x.to(u.m), y.to(u.m)
...
>>> x, y = projectile_motion(
...     5.0 * u.km / u.h, 60 * u.deg, np.linspace(0, 1, 50) * u.s
... )

Technical details & limitations

The most significant limitation of this library is the fact that users must use jpu.numpy functions when interacting with "quantities" with units instead of the jax.numpy interface. This is because JAX does not (yet?) provide a general interface for dispatching of ufuncs on custom array classes. I have played around with the undocumented __jax_array__ interface, but it's not really flexible enough, and it isn't currently compatible with Pytree objects.

So far, only a subset of the numpy/jax.numpy interface is implemented. Pull requests adding broader support (including submodules) would be welcome!

Project details


Download files

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

Source Distribution

jpu-0.0.5.tar.gz (19.0 kB view details)

Uploaded Source

Built Distribution

jpu-0.0.5-py3-none-any.whl (15.9 kB view details)

Uploaded Python 3

File details

Details for the file jpu-0.0.5.tar.gz.

File metadata

  • Download URL: jpu-0.0.5.tar.gz
  • Upload date:
  • Size: 19.0 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/6.1.0 CPython/3.12.8

File hashes

Hashes for jpu-0.0.5.tar.gz
Algorithm Hash digest
SHA256 e0b7e879e5f56dbd2e496813aaae2da0ccb68120aba3f77cbbff0faa1456fc4b
MD5 df5fcd1a91781d279b6ca97060938197
BLAKE2b-256 01090664eb92c03d1e80e2b6fd852c76ed08cbbb9b96859534322906b34c5b6c

See more details on using hashes here.

Provenance

The following attestation bundles were made for jpu-0.0.5.tar.gz:

Publisher: tests.yml on dfm/jpu

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

File details

Details for the file jpu-0.0.5-py3-none-any.whl.

File metadata

  • Download URL: jpu-0.0.5-py3-none-any.whl
  • Upload date:
  • Size: 15.9 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/6.1.0 CPython/3.12.8

File hashes

Hashes for jpu-0.0.5-py3-none-any.whl
Algorithm Hash digest
SHA256 89a4f35f96e8c49fe3be8813e2f3bf122dacf4386fb7471ed97ba1a41e3d31aa
MD5 c7b7a70e4e529f84be6494854cd6dadf
BLAKE2b-256 ffd4ac67e27bd9a95c3839b5808239fa0ae6b925a130a9b89cfb4f0a5bb8b94e

See more details on using hashes here.

Provenance

The following attestation bundles were made for jpu-0.0.5-py3-none-any.whl:

Publisher: tests.yml on dfm/jpu

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

Supported by

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