Motivation

SAIUnit (/saɪ ˈjuːnɪt/) is designed to provide physical units and unit-aware mathematical systems tailored for Scientific AI within JAX. In this context, Scientific AI refers to the use of AI models or tools to advance scientific computations. SAIUnit evolves from our BrainUnit, a unit framework originally developed for brain dynamics modeling, extending its capabilities to support a broader range of scientific computing applications. SAIUnit is committed to providing rigorous and automatic physical unit conversion and analysis system for general AI-driven scientific computing.

Features

Compared to existing unit libraries, such as Quantities and Pint, SAIUnit introduces a rigorous physical unit system specifically designed to support AI computations (e.g., automatic differentiation, just-in-time compilation, and parallelization). Its unique advantages include:

• Integration of over 2,000 commonly used physical units and constants
• Implementation of more than 500 unit-aware mathematical functions
• Deep integration with JAX, providing comprehensive support for modern AI framework features including automatic differentiation (autograd), just-in-time compilation (JIT), vectorization, and parallel computation
• Unit conversion and analysis are performed at compilation time, resulting in zero runtime overhead
• Strict physical unit type checking and dimensional inference system, detecting unit inconsistencies during compilation

graph TD
    A[SAIUnit] --> B[Physical Units]
    A --> C[Mathematical Functions]
    A --> D[JAX Integration]
    B --> B1[2000+ Units]
    B --> B2[Physical Constants]
    C --> C1[500+ Unit-aware Functions]
    D --> D1[Autograd]
    D --> D2[JIT Compilation]
    D --> D3[Vectorization]
    D --> D4[Parallelization]

Loading

We hope these features establish SAIUnit as a reliable physical unit handling solution for general AI-driven scientific computing scenarios.

A quick example:

import saiunit as u

# Define a physical quantity
x = 3.0 * u.meter
x
# [out] 3. * meter

# autograd
f = lambda x: x ** 3
u.autograd.grad(f)(x)
# [out] 27. * meter2 


# JIT
import jax
jax.jit(f)(x)
# [out] 27. * klitre

# vmap
jax.vmap(f)(u.math.arange(0. * u.mV, 10. * u.mV, 1. * u.mV))
# [out]  ArrayImpl([  0.,   1.,   8.,  27.,  64., 125., 216., 343., 512., 729.]) * mvolt3

Multiple-backend support

saiunit is backend-agnostic: a Quantity pairs a unit with an array mantissa, and that mantissa can live on any of the supported array libraries. Every unit-aware operation dispatches to the matching backend, so you can stay in one library end-to-end or convert with a single method call.

Backend	Mantissa	Install	When to use
numpy	numpy.ndarray	core (always installed)	eager CPU, scipy/pandas/sklearn interop
jax	jax.Array	saiunit[jax] (or [cpu]/[cuda12]/[cuda13]/[tpu])	autograd, JIT, vmap, accelerators
cupy	cupy.ndarray	saiunit[cupy]	NVIDIA GPU arrays
torch	torch.Tensor	saiunit[torch]	PyTorch models, torch autograd
dask	dask.array.Array	saiunit[dask]	out-of-core / parallel, lazy compute
ndonnx	ndonnx.Array	saiunit[ndonnx]	symbolic graph for ONNX export

Select or override the backend explicitly with u.using_backend(...) or u.set_default_backend(...); convert with q.to_jax() / q.to_numpy() / q.to_cupy() / q.to_torch() / q.to_dask() / q.to_ndonnx(). Requesting an uninstalled backend raises saiunit.BackendError with the install command — not a bare ImportError. See the Backends documentation for the full story.

Installation

saiunit has been well tested on python>=3.10 and can be installed on Windows, Linux, and MacOS. The core package depends only on NumPy. JAX is optional — install it to enable the saiunit.autograd, saiunit.lax, and saiunit.sparse submodules, the custom exprel primitive, and the "jax" backend.

Install the NumPy-only core:

pip install saiunit --upgrade

Or pull in JAX with the accelerator build that matches your hardware:

pip install -U saiunit[jax]      # plain JAX
pip install -U saiunit[cpu]      # pinned JAX CPU wheels
pip install -U saiunit[cuda12]   # JAX on CUDA 12
pip install -U saiunit[cuda13]   # JAX on CUDA 13
pip install -U saiunit[tpu]      # JAX on TPU

Opt into additional array backends with the matching extra:

pip install -U saiunit[cupy]     # CuPy (NVIDIA GPU)
pip install -U saiunit[torch]    # PyTorch
pip install -U saiunit[dask]     # Dask
pip install -U saiunit[ndonnx]   # ndonnx
pip install -U saiunit[all]      # jax + cupy + torch + dask + ndonnx

Without JAX, the NumPy backend is auto-selected and any access to a JAX-only submodule (saiunit.autograd, saiunit.lax, saiunit.sparse) raises saiunit.BackendError with an install hint. The optional extras are independent and can be combined freely.

To install the latest version from source:

git clone https://github.com/chaobrain/saiunit.git
cd saiunit
pip install -e .

Alternatively, you can install BrainX, which bundles saiunit with other compatible packages for a comprehensive brain modeling ecosystem:

pip install BrainX -U

Documentation

The official documentation is hosted on Read the Docs: https://saiunit.readthedocs.io

Citation

@article{wang2025integrating,
  title={Integrating physical units into high-performance AI-driven scientific computing},
  author={Wang, Chaoming and He, Sichao and Luo, Shouwei and Huan, Yuxiang and Wu, Si},
  journal={Nature Communications},
  volume={16},
  number={1},
  pages={3609},
  year={2025},
  publisher={Nature Publishing Group UK London},
  url={https://doi.org/10.1038/s41467-025-58626-4}
}

Ecosystem

saiunit has been deeply integrated into following diverse projects, such as:

• brainstate: A State-based Transformation System for Program Compilation and Augmentation
• braintaichi: Leveraging Taichi Lang to customize brain dynamics operators
• braintools: The Common Toolbox for Brain Dynamics Programming.
• dendritex: Dendritic Modeling in JAX
• pinnx: Physics-Informed Neural Networks for Scientific Machine Learning in JAX.

Other unofficial projects include:

• diffrax: Numerical differential equation solvers in JAX.
• jax-md: Differentiable Molecular Dynamics in JAX
• Catalax: JAX-based framework to model biological systems
• ...

Acknowledgement

The initial version of the project benefited a lot from the following projects:

• astropy.units: physical units in astropy.
• brian2.units: physical units in brian2.