unxt: A Python package for unit-aware computing with JAX

TL;DR

unxt is a Python package that integrates unit-aware computing into JAX, enabling high-performance scientific computations with physical units by extending the quax framework and leveraging astropy.units.

Contribution

It introduces unxt, a novel extension of quax, to support unit-aware array computations in JAX, enhancing scientific computing capabilities.

Findings

Seamless integration of units into JAX computations

Enhanced support for physical units in high-performance computing

Improved accuracy and clarity in scientific code

Abstract

unxt is a Python package for unit-aware computing with JAX. unxt is built on top of quax, which provides a framework for building array-like objects that can be used with JAX. unxt extends quax to provide support for unit-aware computing using the astropy.units package as a units backend. unxt provides seamless integration of physical units into high performance numerical computations, significantly enhancing the capabilities of JAX for scientific applications.