Fig. 2: JAX-compatible physical units and unit-aware mathematical functions.
From: Integrating physical units into high-performance AI-driven scientific computing
a Categories of physical units supported by SAIUnit, spanning over 2000 commonly used units in fields such as physics, chemistry, astronomy, geoscience, biology, psychology, materials science, and engineering. b Unit-aware mathematical functions in SAIUnit follow NumPy’s array programming conventions16. Supported operations include but not limited to arithmetic, statistics, bitwise manipulation, linear algebra, array manipulation, searching, sorting, broadcasting, aggregation, and matrix operations. c All physical units and mathematical functions in SAIUnit are fully compatible with JAX transformations, including jax.grad for automatic differentiation, jax.vmap for vectorization, jax.pmap for parallelization, and jax.jit for JIT compilation.
