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

pip install quaxed

Note 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!

To run tests locally, you'll need to install nox with nox-uv:

uv tool install --with nox-uv nox

Then you can run the test suite:

nox -s test

Or run linting:

nox -s lint

The quaxed.numpy module uses automatically generated type stubs. These are typically generated during package installation, but if you need to regenerate them (e.g., after updating JAX), pass the --remake-stubs flag:

nox -s lint -- --remake-stubs
nox -s test -- --remake-stubs

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.