galacticdynamics / quaxed

Pre-`quaxify`d libraries for use in multiple dispatch over abstract array types in JAX.

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

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!

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.

quaxed's Issues

More structured approach to adding the numpy namespace

rather than the extremely inefficient piecemeal

Add pytest-github-actions-annotate-failures to CI

https://github.com/pytest-dev/pytest-github-actions-annotate-failures

Wrap `jax.experimental.array_api`

Instead of rolling my own compatability. This will be much easier.

Change repo name

The real thing this repo adds is the quax wrapping. The name should better reflect this.

Rename repo (and PyPI) to `quaxed`

Then we can have quaxed.array_api, quaxed.numpy, etc to cover the full namespace in one area.

jit where appropriate

Figure out which functions can be jitted (with inline=True to let them be inlined in containing jit).

Add tests

• unit tests
• test whether something can be more directly imported from jax.experimental.array_api. In particular, this repo adds jaxtyping + beartype runtime-checking support, which relies on type hints that are not in jax.experimental.array_api.