unxt

unxt

Unitful quantities and calculations in JAX.

unxt gives you arrays that carry their physical units, and keeps them carrying those units through everything JAX does: JIT compilation, vectorization, automatic differentiation, GPU and TPU. There are no unit-aware re-exports of JAX to learn. Your existing JAX code works, with one decorator.

>>> import unxt as u

>>> v = u.Q(25.0, "m/s")
>>> t = u.Q(3.0, "s")
>>> v * t
Quantity(Array(75., dtype=float32, weak_type=True), unit='m')

Install it with pip install unxt — see How to install unxt for the other options.

---

Where to start

🎓 I want to learn the basics

Tutorials. Guided lessons that build something from nothing and run it through JAX. Start here if unxt is new to you.

Tutorials

🔧 I need to get something done

How-to guides. Directions for one specific task: converting units, using JAX functions, controlling display, migrating from v1, making it fast.

How-to guides

📇 I need to look something up

Reference. What every class, function and option is and does, plus the generated API documentation and the glossary.

Reference

💡 I want to understand how it works

Discussion. Why unxt is built the way it is, and where units and JAX surprise people.

Discussion

Looking for ParametricQuantity, unitful linear algebra, hypothesis strategies, or interop with gala, matplotlib or xarray? Those live in the separate unxts.* packages, listed under Packages in the sidebar.

Upgrading from v1? See How to migrate to v2.

---

For citation information and the surrounding ecosystem, see About.