Thalassa: Transforming Symbolic PDEs into Tensor-Based Solvers Running on ML Accelerators
Michail Boulasikis, Flavius Gruian, Robert-Zoltán Szász · 2025
We introduce Thalassa, a framework designed to convert nonlinear systems of partial differential equations (PDEs) with a time-like component into tensor programs that solve these equations. These programs can run on GPUs as well as machine learning (ML) acceleration hardware, enabling scientific computing fields such as computational fluid dynamics, astrophysics, mechanics and biology to utilize any of these resources. Thalassa accepts as input a PDE system expressed in symbolic form and a discretization strategy for the system’s derivatives, in the form of an explicit finite difference method (FDM). With these inputs, Thalassa generates a PyTorch program that solves the system of PDEs numerically using the specified FDM. The generated solver is described as a neural network via a subset of PyTorch operations commonly supported in ML accelerators. However, unlike other ML-based PDE solvers that learn the target PDE system via backpropagation, our approach does not involve training the network. Instead, we directly implement the given numerical method by fixing the weights of the network, making inference correspond to integration. Our results show that a wide variety of PDE systems can be solved by our generated solvers. In addition, the solvers can be generated for either a CPU, GPU or ML accelerators only by changing one parameter in the framework. Finally, we explore the performance bottlenecks of our solvers and their efficiency in using the targeted hardware, identifying areas for optimization and hardware-software co-design as future research directions.