The field of machine learning often uses Python and its PyTorch library by default for both education and real-world industry projects, due to the perceived friendliness of the language and its ability to run on GPU hardware. However, the inability to track tensor dimensions and data types has resulted in most practical code being littered with comments that hopefully explain the current shape of the calculation. There should be a better way.

In this presentation, Jan Ypma first gives an introduction to tensors and the role they play in machine learning. Then, he presents an approach that uses the Scala type system to track all dimensions, the data type, and the device of a PyTorch tensor. He shows how concepts like matrix multiplication and broadcasting rules are applied in a type-safe manner, and how it is possible to refer to specific dimensions by their type, rather than by their index, all while preserving the same GPU performance that PyTorch normally offers.

A typical “transformer” neural network architecture is shown, with an exploration of how having strong types tremendously improves code discoverability. Nice examples also demonstrate how Scala-specific syntax can further improve the legibility of neural network code.