unxt

Unitful Quantities in JAX

Unxt is unitful quantities and calculations in JAX, built on Equinox and Quax.

Unxt supports JAX's compelling features:

JIT compilation (jit)
vectorization (vmap, etc.)
auto-differentiation (grad, jacobian, hessian)
GPU/TPU/multi-host acceleration

And best of all, unxt doesn't force you to use special unit-compatible re-exports of JAX libraries. You can use unxt with existing JAX code, and with quax's simple decorator, JAX will work with unxt.Quantity.

Installation

pip install unxt

using uv

uv add unxt

from source, using pip

pip install git+https://github.com/GalacticDynamics/unxt.git

building from source

cd /path/to/parent
git clone https://github.com/GalacticDynamics/unxt.git
cd unxt
pip install -e .  # editable mode

Documentation

For full documentation, including installation instructions, tutorials, and API reference, please see the unxt docs. This README provides a brief overview and some quick examples.

Dimensions

Dimensions represent the physical type of a quantity, such as length, time, or mass.

>>> import unxt as u

Create dimensions from strings:

>>> u.dimension("length")
PhysicalType('length')

Dimensions support mathematical expressions:

>>> u.dimension("length / time")
PhysicalType({'speed', 'velocity'})

Multi-word dimension names require parentheses in expressions:

>>> u.dimension("(amount of substance) / (time)")
PhysicalType('catalytic activity')

Units

Units specify the scale and dimension of measurements.

>>> meter = u.unit("m")
>>> meter
Unit("m")

Units can be combined, either inside the expression string or by arithmetic on Unit objects:

>>> u.unit("km/h")  # in the expression
Unit("km / h")

>>> u.unit("km") / u.unit("h")  # via arithmetic
Unit("km / h")

Get the dimension of a unit:

>>> u.dimension_of(meter)
PhysicalType('length')

Unit Systems

Unit systems define consistent sets of base units for specific domains. unxt provides built-in unit systems and tools for creating custom ones.

Built-in Unit Systems

>>> u.unitsystem("si")  # SI (International System of Units)
unitsystem(m, kg, s, mol, A, K, cd, rad)

>>> u.unitsystem("cgs")  # CGS (centimeter-gram-second)
unitsystem(cm, g, s, dyn, erg, Ba, P, St, rad)

>>> u.unitsystem("galactic")  # galactic (astrophysics)
unitsystem(kpc, Myr, solMass, rad)

Composing Units from a Unit System

Once you have a unit system, you can get units for any physical dimension by indexing the system:

>>> usys = u.unitsystem("si")
>>> usys["length"]
Unit("m")

Custom Unit Systems

Create custom unit systems by specifying base units:

>>> custom_usys = u.unitsystem("km", "h", "tonne", "degree")
>>> custom_usys
unitsystem(km, h, t, deg)

Derived units are then available by dimension:

>>> custom_usys["velocity"]
Unit("km / h")

Dynamical Unit Systems

For domains like gravitational dynamics, use dynamical unit systems where $G = 1$. Specify only 2 of (length, time, mass); the third is computed to make $G = 1$.

>>> from unxt.unitsystems import DynamicalSimUSysFlag

>>> dyn_usys = u.unitsystem(DynamicalSimUSysFlag, "kpc", "Myr")
>>> dyn_usys
LengthMassTimeUnitSystem(length=Unit("kpc"),
                         mass=Unit("1.49828e+10 kpc3 s2 kg / (Myr2 m3)"), time=Unit("Myr"))

The mass unit is the derived one — an exact composite expression, not a rounded label:

>>> dyn_usys["mass"]
Unit("1.49828e+10 kpc3 s2 kg / (Myr2 m3)")

Quantities

Quantities combine values with units, providing type-safe unitful arithmetic.

Quantity (u.Q) is the lightweight, non-parametric default: a single class — and a single JAX pytree type — for all physical dimensions. ParametricQuantity (up.PQ) adds runtime dimension checking and dimension-specific plum dispatch by encoding each dimension in its own on-the-fly class (and pytree type), which grows the type/dispatch surface and adds per-construction overhead. (This is not about jax.jit cache misses: the unit is static, so a jitted function specializes per unit with either class — that part is inherent.) See the Quantity guide for full details; upgrading from an earlier version? See the migration guide.

Basic Quantities

>>> import jax.numpy as jnp

>>> x = u.Q(jnp.arange(1, 5, dtype=float), "km")
>>> x
Quantity(Array([1., 2., 3., 4.], dtype=float32...), unit='km')

The constituent value and unit are accessible as attributes:

>>> x.value
Array([1., 2., 3., 4.], dtype=float32...)

>>> x.unit
Unit("km")

Quantity objects obey the rules of unitful arithmetic — addition, subtraction, multiplication, division, and exponentiation:

>>> x + x
Quantity(Array([2., 4., 6., 8.], dtype=float32...), unit='km')

>>> 2 * x
Quantity(Array([2., 4., 6., 8.], dtype=float32...), unit='km')

>>> y = u.Q(jnp.arange(4, 8, dtype=float), "yr")
>>> x / y
Quantity(
    Array([0.25     , 0.4      , 0.5      , 0.5714286], dtype=float32...), unit='km / yr'
)

>>> x**2
Quantity(Array([ 1.,  4.,  9., 16.], dtype=float32...), unit='km2')

Operations are unit-checked, so mixing incompatible dimensions raises:

>>> try:
...     x + y
... except Exception as e:
...     print(e)
'yr' (time) and 'km' (length) are not convertible

Quantities can be converted to different units, by function or by method:

>>> u.uconvert("m", x)  # via function
Quantity(Array([1000., 2000., 3000., 4000.], dtype=float32...), unit='m')

>>> x.uconvert("m")  # via method
Quantity(Array([1000., 2000., 3000., 4000.], dtype=float32...), unit='m')

ParametricQuantity — from the separate unxts.parametric package (pip install unxts.parametric, imported below as up) — adds runtime dimension checking on construction. Use up.PQ["length"] to create a parametric type that raises if the unit's physical type does not match:

>>> import unxts.parametric as up

>>> LengthQuantity = up.PQ["length"]
>>> LengthQuantity(2, "km")
ParametricQuantity(Array(2, dtype=int32...), unit='km')

A unit whose physical type does not match the parameter is rejected:

>>> try:
...     LengthQuantity(2, "s")
... except ValueError as e:
...     print(e)
Physical type mismatch.

By contrast, the default u.Q["length"] accepts the subscript but does not check dimensions — it silently builds a Quantity with the mismatched unit:

>>> u.Q["length"](2, "s")
Quantity(Array(2, dtype=int32...), unit='s')

Use up.PQ["length"] when you need the runtime guard. See the unxts.parametric guide for the full API.

Quantity

Quantity (aliased as u.Q) is the default lightweight class. It does not do runtime dimension checking on construction, which makes it the fastest option for performance-critical code:

>>> bq = u.quantity.Quantity(jnp.array([1.0, 2.0, 3.0]), "m")
>>> bq
Quantity(Array([1., 2., 3.], dtype=float32...), unit='m')

>>> bq * 2
Quantity(Array([2., 4., 6.], dtype=float32...), unit='m')

Angle

Angle is a specialized quantity with wrapping support for angular values:

>>> theta = u.Angle(jnp.array([0, 90, 180, 270, 360]), "deg")
>>> theta
Angle(Array([  0,  90, 180, 270, 360], dtype=int32...), unit='deg')

Angles can optionally be wrapped into a specified range:

>>> angle = u.Angle(jnp.array([370, -10]), "deg")
>>> angle.wrap_to(u.Q(0, "deg"), u.Q(360, "deg"))
Angle(Array([ 10, 350], dtype=int32...), unit='deg')

StaticQuantity

For static configuration values (e.g., JAX static arguments), use StaticQuantity, which stores NumPy values and rejects JAX arrays:

>>> import numpy as np
>>> from functools import partial
>>> import jax

>>> cfg = u.StaticQuantity(np.array([1.0, 2.0]), "m")

>>> @partial(jax.jit, static_argnames=("q",))
... def add(x, q):
...     return x + jnp.asarray(q.value)

>>> add(1.0, cfg)
Array([2., 3.], dtype=float32...)

StaticValue

If you want a Quantity that keeps a static value but still participates in regular arithmetic, wrap the value with StaticValue. Arithmetic behaves like the wrapped array, and StaticValue + StaticValue returns a StaticValue. Equality between two StaticValues (== / !=) returns a scalar bool — which is what makes a StaticValue-backed quantity hashable and usable as a jax.jit static argument. Ordering (<, <=, >, >=), and == / != against a raw array, return element-wise NumPy boolean arrays:

>>> sv = u.quantity.StaticValue(np.array([1.0, 2.0]))
>>> q_static = u.Q(sv, "m")
>>> q = u.Q(jnp.array([3.0, 4.0]), "m")

>>> q_static + q
Quantity(Array([4., 6.], dtype=float32...), unit='m')

Equality between two StaticValues is a scalar bool:

>>> sv2 = u.quantity.StaticValue(np.array([2.0, 1.0]))
>>> sv == sv2
False
>>> sv == u.quantity.StaticValue(np.array([1.0, 2.0]))
True

Ordering, and equality against a raw array, are element-wise NumPy boolean arrays:

>>> sv < sv2
array([ True, False])

>>> sv == np.array([1.0, 2.0])
array([ True,  True])

JAX Integration

unxt is built on quax, which enables custom array-ish objects in JAX. For convenience we use the quaxed library, which is just a quax.quaxify wrapper around jax to avoid boilerplate code.

[!NOTE]

Using quaxed is optional. You can directly use quaxify, and even apply it to the top-level function instead of individual functions.

Using the x quantity from the earlier examples:

>>> from quaxed import grad, vmap
>>> import quaxed.numpy as qnp

>>> qnp.square(x)
Quantity(Array([ 1.,  4.,  9., 16.], dtype=float32...), unit='km2')

>>> qnp.power(x, 3)
Quantity(Array([ 1.,  8., 27., 64.], dtype=float32...), unit='km3')

>>> vmap(grad(lambda x: x**3))(x)
Quantity(Array([ 3., 12., 27., 48.], dtype=float32...), unit='km2')

See the documentation for more examples and details of JIT and AD

Citation

If you found this library to be useful and want to support the development and maintenance of lower-level code libraries for the scientific community, please consider citing this work.

Contributing and Development

We welcome contributions! Contributions are how open source projects improve and grow.

To contribute to unxt, please fork the repository, make a development branch, develop on that branch, then open a pull request from the branch in your fork to main.

To report bugs, request features, or suggest other ideas, please open an issue.

For more information, see CONTRIBUTING.md.

Metadata

Release files for unxt 2.0.4

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for unxt 2.0.4
File	Size	Uploaded
unxt-2.0.4.tar.gz	725.4 kB	Sep 4, 2026	Details

Built distribution (wheel)

Table of built distributions (wheels) for unxt 2.0.4
File	Interpreter	ABI	Platform	Reset
unxt-2.0.4-py3-none-any.whl 119.9 kB Sep 4, 2026	Python 3	none	any	Details

Total release size: 845.4 kB

Release files / unxt-2.0.4.tar.gz

Download URL	unxt-2.0.4.tar.gz
Size	725.4 kB
Tags	Source
SHA-256 checksum How to use checksums	`fb226465242c9465045dfee71f09757afafb312494a22cc2c9b82b28842bc572`
BLAKE2b-256 checksum How to use checksums	`0adfc6809f0be3fef6793a68d5fd6ee0a5e5174dcc40f3961c39367d6de329a6`
Upload date	Sep 4, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	Yes
Uploaded via	`twine/7.0.0 CPython/3.13.14`

Provenance

Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.

Identity

Publishing platform: GitHub Actions
Publishing repository: github.com/GalacticDynamics/unxt
Publishing commit: github.com/GalacticDynamics/unxt/tree/dddcf650e3684ba183ca7bb388ebc2369ff34aea

Workflow

Publishing configuration: github.com/GalacticDynamics/unxt/blob/dddcf650e3684ba183ca7bb388ebc2369ff34aea/.github/workflows/cd-publish-unxt.yml
Publishing logs: github.com/GalacticDynamics/unxt/actions/runs/33851400834/attempts/1

Artifact

Subject: unxt-2.0.4.tar.gz
SHA-256 checksum: fb226465242c9465045dfee71f09757afafb312494a22cc2c9b82b28842bc572 Verifying attestations

Release files / unxt-2.0.4-py3-none-any.whl

Download URL	unxt-2.0.4-py3-none-any.whl
Size	119.9 kB
Tags	Python 3
SHA-256 checksum How to use checksums	`8443041705ca3a04ee78b9c6e48d4f817f1cc54d3a7862203b5da0ee79cf14fd`
BLAKE2b-256 checksum How to use checksums	`4e1a33374c3793f6baf0534833a246820951aae2fa6873567586101e8af8067e`
Upload date	Sep 4, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	Yes
Uploaded via	`twine/7.0.0 CPython/3.13.14`

Provenance

Identity

Publishing platform: GitHub Actions
Publishing repository: github.com/GalacticDynamics/unxt
Publishing commit: github.com/GalacticDynamics/unxt/tree/dddcf650e3684ba183ca7bb388ebc2369ff34aea

Workflow

Publishing configuration: github.com/GalacticDynamics/unxt/blob/dddcf650e3684ba183ca7bb388ebc2369ff34aea/.github/workflows/cd-publish-unxt.yml
Publishing logs: github.com/GalacticDynamics/unxt/actions/runs/33851400834/attempts/1

Artifact

Subject: unxt-2.0.4-py3-none-any.whl
SHA-256 checksum: 8443041705ca3a04ee78b9c6e48d4f817f1cc54d3a7862203b5da0ee79cf14fd Verifying attestations