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

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.

arXiv comment4 pages, software

References in corpus (4)

• Array Programming with NumPy
• Astropy: A Community Python Package for Astronomy
• The Astropy Project: Sustaining and Growing a Community-oriented Open-source Project and the Latest Major Release (v5.0) of the Core Package
• unyt: Handle, manipulate, and convert data with units in Python