DimTensor: A Unit-Aware Tensor Computation Framework for Physics and Machine Learning

This post introduces the DimTensor open-source project, a unit-aware tensor computation framework for physics and machine learning. Its core goal is to enable real-time dimension error detection in tensor operations by integrating physical unit checks into the tensor type system. Key features include treating units as intrinsic to the tensor type system, dual compile-time/runtime checks for unit mismatches, and seamless integration with PyTorch/JAX. DimTensor aims to reduce unit confusion errors and improve safety/reliability in scientific computing and ML workflows.

In scientific computing and ML, tensors are core data structures, but traditional libraries (NumPy, PyTorch) lack native physical unit support. This leads to accidental mixing of units (e.g., meters and feet) without warnings, causing incorrect results. DimTensor was created to address this pain point by adding unit-aware mechanisms to tensor operations, ensuring high performance while detecting dimension mismatches in real time.

DimTensor's design follows three principles: 1. Units as part of the type system: Physical units are intrinsic tensor properties (not metadata) with clear dimensions (length, mass, time) tracked during operations. 2. Dual checks: Compile-time error capture where possible, plus runtime checks for reliability. 3. Seamless integration: Extends existing libraries (PyTorch/JAX) for gradual adoption without large refactoring.

DimTensor's implementation includes: 1. Dimension algebra system: Units are combinations of basic dimensions (L, M, T) with exponents (e.g., speed = L·T⁻¹). Operations validate legality (add requires same units; multiply follows dimension rules). 2. Tensor operation tracking: Tracks unit info during neural network forward/backward passes to locate error sources in complex graphs. 3. Detailed error reporting: Provides tensor details, expected/actual units, and code location to shorten debugging time.

Application Scenarios: - Physical simulation/engineering: Avoids unit confusion in aerospace/automotive engineering (critical for safety). - Scientific ML: Ensures model outputs align with physical constraints in physics-integrated neural networks. - Finance/econometrics: Supports dimension-consistent pricing and risk models. Comparison: - vs Pint: DimTensor focuses on tensor operations/auto-differentiation (GPU support), unlike Pint's limited tensor/gradient support. - vs Unitful.jl: DimTensor targets Python ecosystem with PyTorch/JAX interoperability, while Unitful.jl uses Julia's type system for compile-time checks.

Usage: Typical workflow: define unit systems → create unit-aware tensors → execute operations → check results (docs have examples from simple physics to complex neural networks). Best Practices: Annotate units during data loading, add unit type hints for function parameters, insert unit assertions at key nodes. Community: DimTensor is actively developed (core features available; API iterated via user feedback). Contributions welcome from physics, engineering, ML fields (to improve unit libraries, expand framework support, optimize performance).

Future Plans: Support more physical constants/derived units, smarter unit inference, interfaces with physical simulation software, and visualization tools for unit flow in computation graphs. Conclusion: DimTensor represents a trend toward safer scientific computing tools. By automating unit checks, it reduces human error and improves efficiency. It is a valuable tool for projects involving physical quantity computations.