Quaxed wraps jax libraries (using
quax) to enable using those libraries with
custom array-ish objects, not only jax arrays.
pip install quaxedTo 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.