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

The paper introduces **unxt**, a Python package built on the **quax** framework that integrates **astropy.units** with **JAX** to enable seamless, high-performance, unit-aware scientific computing.

The Gist

Imagine you are a chef in a high-speed kitchen. You have a super-fast robot assistant (let's call it JAX) that can chop, mix, and cook ingredients faster than any human ever could. It can even predict exactly how the flavors will change if you tweak a recipe, and it does all this on different types of stoves (CPUs, GPUs, TPUs) without breaking a sweat.

However, there's a catch: JAX only understands raw numbers. It doesn't know the difference between "5 cups of flour" and "5 pounds of sugar." If you tell it to add 5 cups to 5 pounds, it will just blindly add the numbers together to get "10" and serve you a disaster. In the real world, mixing up units (like the famous Mars Climate Orbiter crash, where one team used metric and another used imperial) can lead to catastrophic failures.

Enter unxt.

What is unxt?

Think of unxt as a smart translator and safety guard that sits between you and the super-fast robot.

1. The Translator: It takes your "real-world" instructions (like "5 meters per second") and translates them into a format the robot understands, while keeping a tag on every single number that says, "Hey, this is a length!" or "This is time!"
2. The Safety Guard: If you try to do something silly, like adding "5 meters" to "5 seconds," the translator stops you immediately and says, "Whoa, you can't mix those up!" This prevents the robot from making a mistake that could ruin your experiment or launch a rocket into the wrong planet.

How does it work?

The paper explains that unxt is built on top of a framework called quax.

• Quax is like a universal adapter that lets the robot (JAX) talk to special, custom objects.
• unxt uses this adapter to create a special kind of "smart number" (called a Quantity) that knows it has units attached to it.
• It also uses the Astropy library (a giant, trusted encyclopedia of units used by astronomers) as its brain for knowing what a "kilogram" or a "light-year" actually means.

Why do we need it?

Before unxt, scientists who wanted to use the super-fast robot (JAX) had to choose:

• Option A: Use JAX for speed, but lose the ability to check units (risky!).
• Option B: Use a library like Astropy or Pint to check units, but lose the speed and power of JAX (slow!).

unxt solves this by letting you have your cake and eat it too. It lets you write code that is:

• Fast: It runs on JAX's high-speed engine.
• Safe: It automatically checks that you aren't mixing apples and oranges.
• Easy: It feels familiar to scientists who already use Astropy, so they don't have to learn a whole new language.

The Bottom Line

unxt is a bridge. It connects the messy, complex world of physical measurements (meters, seconds, kilograms) with the clean, lightning-fast world of modern computer science. It ensures that when scientists run complex simulations—whether for weather forecasting, astrophysics, or engineering—the math is not only fast but also correct.

In short: It's the safety net that lets the super-fast robot cook without burning the kitchen down.

Technical Summary

Based on the provided paper, here is a detailed technical summary of unxt, covering the problem statement, methodology, key contributions, results, and significance.

1. Problem Statement

The paper identifies a critical gap in the high-performance scientific computing ecosystem: the lack of native support for unit-aware computing within JAX.

• JAX Limitations: While JAX excels in automatic differentiation, just-in-time (JIT) compilation, and hardware acceleration (CPUs, GPUs, TPUs), it operates primarily on "pure" arrays. It lacks the ability to define custom array-like objects that carry physical units (e.g., meters, seconds) without breaking the JAX execution model.
• Astropy Limitations: Astropy, the standard library for astronomy and physics, offers robust unit handling via astropy.units (used in over 10,000 citations). However, astropy.units is built on NumPy and cannot be directly extended to work with JAX's functional programming style or its JIT compilation requirements.
• Consequence: This disconnect forces researchers to choose between the performance and differentiation capabilities of JAX and the safety and correctness of unit-aware calculations. This trade-off increases the risk of catastrophic errors in scientific simulations (e.g., unit mismatches similar to the Mars Climate Orbiter incident).

2. Methodology

The authors developed unxt as a Python package to bridge this gap by creating a seamless integration between JAX and unit management. The methodology relies on a layered architecture:

• Foundation on Quax: unxt is built on top of quax (Kidger, 2023), a framework designed to create array-like objects compatible with JAX. Quax allows for the definition of custom types that JAX can recognize and differentiate.
• Backend Integration: unxt utilizes astropy.units as its backend for unit conversion and physical logic. This ensures compatibility with the established standards of the scientific community.
• The Quantity Class: The core innovation is the definition of a Quantity class. This class wraps data arrays and attaches unit metadata.
• Multiple Dispatch & Overrides: To make Quantity objects work with JAX, unxt leverages multiple dispatch to provide a comprehensive set of overrides for JAX primitives. This allows JAX functions (like jnp.add, jnp.sin, etc.) to operate transparently on Quantity objects, automatically propagating units through calculations.
• Interface Design: The package offers a function-oriented framework consistent with JAX's style, while maintaining an object-oriented front-end familiar to users of astropy.units. This dual approach lowers the learning curve for existing users while adhering to JAX best practices.

3. Key Contributions

• First JAX-Native Unit System: unxt is the first package to provide a fully functional, unit-aware computing layer specifically optimized for the JAX ecosystem.
• Seamless Interoperability: It enables the use of astropy.units logic within JAX code without requiring users to manually manage unit conversions or worry about underlying JAX interfacing.
• Flexibility in Unit Systems: The package supports both static and dynamic definitions of unit systems, allowing flexibility across various computational environments.
• Extensibility: By leveraging multiple dispatch, unxt is designed to be extensible to other libraries (currently supporting astropy and custom array-like objects), making it a hub for unit-aware interoperability.
• Performance Preservation: The design ensures that unit handling does not compromise JAX's core advantages: automatic differentiation, JIT compilation, and hardware acceleration.

4. Results

• Functional Integration: The paper demonstrates that unxt successfully allows users to perform unit-aware computations effortlessly. Operations involving physical quantities are handled correctly and consistently.
• Error Prevention: By enforcing unit consistency at the code level, the package mitigates the risk of dimensionality errors in scientific calculations.
• Usability: The interface is described as intuitive and easy to use, facilitating the straightforward implementation of units into existing JAX codebases.
• Community Validation: The development team includes core maintainers of astropy.units, ensuring the package aligns with established scientific standards. The work was presented and discussed at the 2024 JAXtronomy workshop, indicating early community engagement.

5. Significance

• Scientific Safety: unxt addresses a fundamental safety issue in computational science by preventing unit-related errors that can lead to significant scientific or engineering failures.
• Enabling Advanced Physics: It unlocks the full potential of JAX for scientific applications that require rigorous physical constraints, such as astrophysics, cosmology, and engineering simulations, where unit consistency is non-negotiable.
• Bridging Ecosystems: It effectively bridges the gap between the mature, unit-rich ecosystem of Astropy and the high-performance, differentiable ecosystem of JAX.
• Future-Proofing: As scientific computing moves toward differentiable programming and machine learning on physical systems, unxt provides the necessary infrastructure to ensure these models remain physically grounded and interpretable.

In summary, unxt is a critical infrastructure tool that democratizes high-performance, unit-aware computing in JAX, combining the speed of modern hardware acceleration with the rigor of physical unit management.