MIT researchers extend tensor programming to the continuous world

CTA is akin to creating a large collage by cutting and pasting rectangles from millions of different-colored plain paper (for example, the art of Piet Mondrian, depicted here): the essential information is preserved, but in a form the hardware can handle (Credit: ChatGPT).

When the FORTRAN programming language debuted in 1957, it transformed how scientists and engineers programmed computers. Complex calculations could suddenly be expressed in concise, math-like notation using arrays — collections of values that make it easier to describe operations on data. That simple idea evolved into today’s “tensors,” which power many of the world’s most advanced AI and scientific computing systems through modern frameworks like NumPy and PyTorch.

Traditional tensor programming assumes that all data live on an integer grid — for instance, values stored at coordinates (1, 2) or (5, 7). This assumption has enabled decades of optimization in both hardware and software. Tensors, with data on integer grids, are now seemingly everywhere in programming, and are foundational to many computing paradigms such as machine learning and scientific computing. But many real-world datasets don’t fit neatly into this grid. Point clouds in 3D sensing, geometric models in computer graphics, or numerical simulations in physics often live in the continuous world, where coordinates can take any real value. To handle these, researchers have long resorted to complex data structures and lengthy, hard-to-maintain code written separately from mainstream tensor frameworks.

Now, researchers at MIT’s Computer Science and Artificial Intelligence Laboratory (CSAIL) have found a way to bring those continuous datasets into the familiar tensor world. Their new framework, called the continuous tensor abstraction (CTA), allows data to be stored and accessed at real-number coordinates — for example, enabling programmers to write `A[3.14]` instead of just `A[3]`. It also introduces a new language called continuous Einsums, a generalization of the Einstein summation notation widely used in traditional tensor programming and originally introduced by Albert Einstein, to express computations between continuous tensors in concise, math-like form. 

Representing continuous data as an array seems impossible — after all, real numbers are infinite and arrays have a finite size. The MIT team overcame this challenge with a clever concept called piecewise-constant tensors, which divide continuous space into manageable regions that have the same value. It’s akin to creating a large collage by cutting and pasting rectangles from millions of different-colored plain paper (for example, the art of Piet Mondrian, depicted in the banner): the essential information is preserved, but in a form the hardware can handle.

This approach makes it possible to express previously cumbersome algorithms in a single, compact line of tensor code. Tasks that once required thousands of lines of specialized logic — from analyzing 3D LiDAR scans to modeling fluid flow or simulating physical systems — can now be written in a familiar, high-level mathematical Einsum language and executed efficiently on modern accelerators.

“Programs that took 2,000 lines of code to write can be done in one line with our language,” says MIT’s Thomas and Gerd Perkins Professor of electrical engineering and computer science and CSAIL principal investigator Saman Amarasinghe, a senior author of the study. “For any task involving continuous data, our work is showing that tensor programming can help you represent such programs in more diverse ways than scientists previously thought.” Another senior author, MIT Professor Joel Emer — a CSAIL principal investigator and Senior Distinguished Research Scientist at NVIDIA — adds, “Continuous Einsums intuitively work the same as the increasingly popular Einsums on tensors defined on regular grids, making the new extension easy for programmers to adopt.”

Making the case for continuous tensors

In one case study, CSAIL researchers found that CTA could help users search 2D spaces, such as those you’d find in geographical information systems (GIS) like Google Maps. The team used two different querying methods: box search, where a system retrieves all points within a square area specified by users, and radius search, where it finds datapoints located within a circle users outline. Across these geospatial tests, CTA generated search programs with 62 times fewer lines of code than the Python tool Shapely, which helps users survey and manipulate shapes. On radius searches, CSAIL’s system was also about nine times faster.

CTA showed even more staggering improvements when programming machine-learning algorithms to analyze patterns across 3D point cloud searches. While implementing the “Kernel Points Convolution” requires over 2,300 lines of code, CSAIL’s approach only needed 23 lines to code the same task — in other words, 101 times more concise.

The system was also tested on its ability to locate features at particular chromosome regions in the human genome. CTA’s craftiness helped execute these search queries, generating code that was 18 times shorter while being slightly faster on average than three comparable baselines.

Another case study showed how useful CTA could be for 3D deep learning tasks, which help robots see their environments before picking up objects. In particular, their system excelled at calculating datapoints within a neural radiance field (NeRF), a deep learning approach that constructs a 3D scene using 2D images. It was nearly twice as fast as a comparable PyTorch tool at this task, while writing roughly 70 lines of code less.

Each experiment revealed how tensors can accelerate how we analyze a variety of complex applications. “Previously, the tensor world and the non-tensor world have largely evolved in isolation, developing their own algorithms and data structures independently,” says lead author Jaeyeon Won SM ‘23, who is an MIT PhD student at CSAIL co-advised by Emer and Amarasinghe. “Using our language, we discovered that many geometric applications can be expressed concisely and precisely in continuous Einsums — while benefiting from the same performance optimizations used in sparse tensor compilers.”

“We believe this serves as a bridge between the two worlds,” adds Won. “Moreover, we aim to explore how optimizations from the non-tensor world — for example, whether ray-tracing cores could accelerate tensor programs — can be reinterpreted and applied through the lens of continuous tensor notation.”

CTA presents an exciting avenue for tensor programming by expanding its capabilities to work in continuous, infinite spaces. Won, Emer, and Amarasinghe’s next step may be to explore even more complex data structures in this realm, which use variables instead of constants to represent the values of certain points within a region. In turn, they could experiment with more complex applications in deep learning and computer graphics, providing a deeper examination of intricate worlds and scientific visualizations.

Won, Emer, and Amarasinghe wrote the paper with Georgia Institute of Technology Assistant Professor Willow Ahrens SM ‘19 PhD ‘24 and CSAIL affiliates: PhD student Teodoro Fields Collin SM ‘25. Their work was supported, in part, by Intel, the National Science Foundation, DARPA, and the U.S. Department of Energy’s Predictive Science Academic Alliance Program (PSAAP). They presented their research at the Object-Oriented Programming, Systems, Languages & Applications conference (OOPSLA) in Singapore in October.