quaxed
Pre-Quaxify'ed JAX
Quaxed wraps jax libraries (using
quax) to enable using those libraries with
custom array-ish objects, not only jax arrays.
Installation
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.
Documentation
Quick Start
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')
Development
We welcome contributions!
Running Tests
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
Citation
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.
Release files for quaxed 0.10.6
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | |
|---|---|---|---|
| quaxed-0.10.6.tar.gz | 172.5 kB | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| quaxed-0.10.6-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 219.0 kB
Release files / quaxed-0.10.6.tar.gz
| Download URL | quaxed-0.10.6.tar.gz |
|---|---|
| Size | 172.5 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
44228c05b6d7d65a6fc320090ebf561f7cef30bd0ce73d8e68bfbf084effcab5
|
|
BLAKE2b-256 checksum How to use checksums |
be663037b535113e21692b04db618f28e51bd8534d42e5e49829ff0008d1be81
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/7.0.0 CPython/3.13.14
|
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
PyPI Publish Attestation
PyPI verified that this artifact, at this checksum, originated from the publisher listed below.
Signed by GitHub Actions, verified by PyPI on Oct 10, 2026.
Transparency logRelease files / quaxed-0.10.6-py3-none-any.whl
| Download URL | quaxed-0.10.6-py3-none-any.whl |
|---|---|
| Size | 46.5 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
82b4ff7c2a9d66f67211741887d3b89a3f6c7c3c8c41ce3490b66d4ce6554626
|
|
BLAKE2b-256 checksum How to use checksums |
68bc375307de632edbd2de26c799933bb6fbbc2d32dfd95e681d0f78227b5c9b
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/7.0.0 CPython/3.13.14
|
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
PyPI Publish Attestation
PyPI verified that this artifact, at this checksum, originated from the publisher listed below.
Signed by GitHub Actions, verified by PyPI on Oct 10, 2026.
Transparency log
