Integrating physical units into high-performance AI-driven scientific computing

Artificial intelligence is revolutionizing scientific research across various disciplines. The foundation of scientific research lies in rigorous scientific computing based on standardized physical units. However, current mainstream high-performance numerical computing libraries for artificial intelligence generally lack native support for physical units, significantly impeding the integration of artificial intelligence methodologies into scientific research. To fill this gap, we introduce SAIUnit, a system designed to seamlessly integrate physical units into scientific artificial intelligence libraries, with a focus on compatibility with JAX. SAIUnit offers a comprehensive library of over 2000 physical units and 500 unit-aware mathematical functions. It is fully compatible with JAX transformations, allowing for automatic differentiation, just-in-time compilation, vectorization, and parallelization while maintaining unit consistency. We demonstrate SAIUnit's applicability and effectiveness across diverse artificial intelligence-driven scientific computing domains, including numerical integration, brain modeling, and physics-informed neural networks. Our results show that by confining unit checking to the compilation phase, SAIUnit enhances the accuracy, reliability, interpretability, and collaborative potential of scientific computations without compromising runtime performance. By bridging the gap between abstract computing frameworks and physical units, SAIUnit represents a crucial step towards more robust and physically grounded artificial intelligence-driven scientific computing.