Skip to main content

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

Project description

quaxed

Pre-Quaxify'ed JAX

PyPI: quaxed PyPI versions: quaxed ReadTheDocs quaxed license

CI status ReadTheDocs codecov ruff ruff pre-commit


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

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

Read The Docs

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 Documentation Status codecov SPEC 0 — Minimum Supported Dependencies pre-commit ruff

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

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.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

quaxed-0.10.5.tar.gz (170.1 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

quaxed-0.10.5-py3-none-any.whl (45.7 kB view details)

Uploaded Python 3

File details

Details for the file quaxed-0.10.5.tar.gz.

File metadata

  • Download URL: quaxed-0.10.5.tar.gz
  • Upload date:
  • Size: 170.1 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/6.1.0 CPython/3.13.7

File hashes

Hashes for quaxed-0.10.5.tar.gz
Algorithm Hash digest
SHA256 97c8b96b46ae832afb83f9e845397738dcbacd61c165bbe9f6542b3bf6892593
MD5 0993fac8fa5d7dd9b2e694ec1a22668e
BLAKE2b-256 8c9f7fdd817ee42e53f7249eecc2bc6f93156a8846d05e2d69e5540a63630747

See more details on using hashes here.

Provenance

The following attestation bundles were made for quaxed-0.10.5.tar.gz:

Publisher: cd.yml on GalacticDynamics/quaxed

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file quaxed-0.10.5-py3-none-any.whl.

File metadata

  • Download URL: quaxed-0.10.5-py3-none-any.whl
  • Upload date:
  • Size: 45.7 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/6.1.0 CPython/3.13.7

File hashes

Hashes for quaxed-0.10.5-py3-none-any.whl
Algorithm Hash digest
SHA256 f5db02cbed3809b7ec349c287fc492ecef226c1ca7063118a7722b1e7a98284c
MD5 a3253dde5df5e812ed24ce601b9f9905
BLAKE2b-256 c27169e5f77bebbac8db7315e26dbacdf21655c6a8992974276d4dff37c2a64a

See more details on using hashes here.

Provenance

The following attestation bundles were made for quaxed-0.10.5-py3-none-any.whl:

Publisher: cd.yml on GalacticDynamics/quaxed

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page