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 rigorious 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

Installation

saiunit has been well tested on python>=3.9 + jax>=0.4.30 environments, and can be installed on Windows, Linux, and MacOS.

You can install saiunit via pip:

pip install saiunit --upgrade

which should install in about 1 minute. If you want to install the latest version from the source, you can clone the repository and install it:

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

Documentation

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

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
• ...