Quaxed wraps jax libraries (using
quax) to enable using those libraries with
custom array-ish objects, not only jax arrays.
pip install quaxedNote that quaxed builds a
typing stub file
during the build process to provide type hints for some quaxified libraries
(e.g., quaxed.numpy). This enables quaxed functions to be typed correctly
when working with custom array-like objects. If you find inconsistencies or
issues with the type hints or stub files, we recommend using
uv to set up your development environment and
pin the versions of quaxed and its dependencies.
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.