Statically Typed Tensor Programming with Semantic Dimension Names

In mainstream programming languages, types can state what values mean and constrain the operations permitted on them, making code easier to comprehend and enabling automated verification of logical consistency. However, common numerical and machine learning libraries forgo these benefits by representing all multidimensional data with generic tensor types. We present DimWit, a numerical and machine learning library in Scala 3 that encodes tensor dimension names as programmer-defined types, such as Batch, allowing tensor types like Tensor2[Batch, Feature, Float32] for a batch of feature vectors, and extending the benefits of static typing to expressing, tracking, and verifying dimension names and data types in tensors and across their operations. DimWit leverages JAX as its computational engine, inheriting hardware acceleration, just-in-time compilation, and automatic differentiation. To explore DimWit's semantic tensor types in practice, we implement diverse applications, ranging from linear regression to variational autoencoders, as well as clustering algorithms and simulations. We highlight excerpts of these implementations as examples, showing how semantic tensor types give component interfaces precise meaning and can make programs more comprehensible. To demonstrate the type-safety benefits, we reconstruct twelve selected axis- and shape-related defects from widely used open-source projects, each of which silently produced incorrect results, and show that DimWit rejects them at compile time. To confirm that the approach adds negligible runtime cost at scale, we train a 124M-parameter GPT-2 from scratch with model and training code written in Scala 3 and execution delegated to JAX, reaching 97% of the throughput of an equivalent JAX implementation in Python. DimWit illustrates how semantic tensor types can make tensor programs more comprehensible and reject axis- and shape-related defects at compile time, without sacrificing runtime performance. Code available at https://anonymous.4open.science/r/dimwit-9234 and https://anonymous.4open.science/r/deepwit-DC06/.