https://github.com/galacticdynamics/quaxed
Pre-quaxed libraries for use in multiple dispatch over abstract array types in JAX.
https://github.com/galacticdynamics/quaxed
array array-api ducktyping jax multiple-dispatch
Last synced: 3 months ago
JSON representation
Pre-quaxed libraries for use in multiple dispatch over abstract array types in JAX.
- Host: GitHub
- URL: https://github.com/galacticdynamics/quaxed
- Owner: GalacticDynamics
- License: bsd-3-clause
- Created: 2023-12-16T04:56:07.000Z (over 1 year ago)
- Default Branch: main
- Last Pushed: 2025-04-09T04:14:56.000Z (3 months ago)
- Last Synced: 2025-04-15T02:06:23.433Z (3 months ago)
- Topics: array, array-api, ducktyping, jax, multiple-dispatch
- Language: Python
- Homepage: https://quaxed.readthedocs.io/en/latest/
- Size: 409 KB
- Stars: 4
- Watchers: 1
- Forks: 1
- Open Issues: 4
-
Metadata Files:
- Readme: README.md
- Contributing: .github/CONTRIBUTING.md
- License: LICENSE
- Code of conduct: CODE_OF_CONDUCT.md
- Security: SECURITY.md
Awesome Lists containing this project
README
quaxed
Pre-Quaxify'ed JAX
---
`Quaxed` wraps [jax](https://jax.readthedocs.io/en/latest/) libraries (using
[`quax`](https://docs.kidger.site/quax/)) to enable using those libraries with
custom array-ish objects, not only jax arrays.
## Installation
[![PyPI version][pypi-version]][pypi-link]
[![PyPI platforms][pypi-platforms]][pypi-link]
```bash
pip install quaxed
```
## Documentation
[](https://unxt.readthedocs.io/en/)
### Quick Start
To understand how `quax` works it's magic, see
[`quax.quaxify`](https://docs.kidger.site/quax/api/quax/#quax.quaxify) and the
[tutorials](https://docs.kidger.site/quax/examples/custom_rules/).
To use this library, it's as simple as:
```pycon
# 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
[![Actions Status][actions-badge]][actions-link]
[![Documentation Status][rtd-badge]][rtd-link]
[![codecov][codecov-badge]][codecov-link]
[![SPEC 0 — Minimum Supported Dependencies][spec0-badge]][spec0-link]
[![pre-commit][pre-commit-badge]][pre-commit-link]
[![ruff][ruff-badge]][ruff-link]
We welcome contributions!
## Citation
[![DOI][zenodo-badge]][zenodo-link]
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.
[actions-badge]: https://github.com/GalacticDynamics/quaxed/workflows/CI/badge.svg
[actions-link]: https://github.com/GalacticDynamics/quaxed/actions
[codecov-badge]: https://codecov.io/gh/GalacticDynamics/quaxed/graph/badge.svg?token=9G19ONVD3U
[codecov-link]: https://codecov.io/gh/GalacticDynamics/quaxed
[pre-commit-badge]: https://img.shields.io/badge/pre--commit-enabled-brightgreen?logo=pre-commit
[pre-commit-link]: https://pre-commit.com
[pypi-link]: https://pypi.org/project/quaxed/
[pypi-platforms]: https://img.shields.io/pypi/pyversions/quaxed
[pypi-version]: https://img.shields.io/pypi/v/quaxed
[rtd-badge]: https://readthedocs.org/projects/quaxed/badge/?version=latest
[rtd-link]: https://quaxed.readthedocs.io/en/latest/?badge=latest
[ruff-badge]: https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/charliermarsh/ruff/main/assets/badge/v2.json
[ruff-link]: https://docs.astral.sh/ruff/
[spec0-badge]: https://img.shields.io/badge/SPEC-0-green?labelColor=%23004811&color=%235CA038
[spec0-link]: https://scientific-python.org/specs/spec-0000/
[zenodo-badge]: https://zenodo.org/badge/732262318.svg
[zenodo-link]: https://zenodo.org/doi/10.5281/zenodo.10850521