Skip to content

Pre-quaxed libraries for use in multiple dispatch over abstract array types in JAX.

License

Notifications You must be signed in to change notification settings

GalacticDynamics/quaxed

Repository files navigation

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

PyPI version PyPI platforms

pip install quaxed

Documentation

Documentation Status

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

Actions Status Codecov

We welcome contributions!

Citation

DOI

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.