Quaxed wraps jax libraries (using quax) to enable using those libraries with custom array-ish objects, not only jax arrays.

pip install quaxed

To understand how quax works it's magic, see quax.quaxify and the tutorials.

To use this library, it's as simple as:

# Import pre-quaxified library
>>> import quaxed.numpy as jnp  # this is quaxify(jax.numpy)

# As an example, let's import an array-ish object
>>> from unxt import Quantity
>>> x = Quantity(2, "km")
>>> jnp.square(w)
Quantity['area'](Array(4, dtype=int64, weak_type=True), unit='km2')

We welcome contributions!

If you found this library to be useful and want to support the development and maintenance of lower-level utility libraries for the scientific community, please consider citing this work.