# Learning, Solving and Optimizing PDEs with TensorGalerkin: an efficient high-performance Galerkin assembly algorithm

arXiv: https://arxiv.org/abs/2602.05052 · HTML: https://mlsanity.com/paper/2602.05052/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 24 of its 47 references that are in the map.
- First posted on arXiv 4 Feb 2026.

## Builds on (most cited first, 24 of 24)

- [Inductive Representation Learning on Large Graphs](https://mlsanity.com/paper/1706.02216/) (2017)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [Implicit Neural Representations with Periodic Activation Functions](https://mlsanity.com/paper/2006.09661/) (2020)
- [Learning Mesh-Based Simulation with Graph Networks](https://mlsanity.com/paper/2010.03409/) (2020)
- [The Deep Ritz method: A deep learning-based numerical algorithm for solving variational problems](https://mlsanity.com/paper/1710.00211/) (2017)
- [Physics-Informed Neural Operator for Learning Partial Differential Equations](https://mlsanity.com/paper/2111.03794/) (2021)
- [Message Passing Neural PDE Solvers](https://mlsanity.com/paper/2202.03376/) (2022)
- [GNOT: A General Neural Operator Transformer for Operator Learning](https://mlsanity.com/paper/2302.14376/) (2023)
- [Transolver: A Fast Transformer Solver for PDEs on General Geometries](https://mlsanity.com/paper/2402.02366/) (2024)
- [Convolutional Neural Operators for robust and accurate learning of PDEs](https://mlsanity.com/paper/2302.01178/) (2023)
- [Poseidon: Efficient Foundation Models for PDEs](https://mlsanity.com/paper/2405.19101/) (2024)
- [Variational Physics-Informed Neural Networks For Solving Partial Differential Equations](https://mlsanity.com/paper/1912.00873/) (2019)
- [Physics-informed graph neural Galerkin networks: A unified framework for solving PDE-governed forward and inverse problems](https://mlsanity.com/paper/2107.12146/) (2021)
- [MultiScale MeshGraphNets](https://mlsanity.com/paper/2210.00612/) (2022)
- [Long-time integration of parametric evolution equations with physics-informed DeepONets](https://mlsanity.com/paper/2106.05384/) (2021)
- [Geometry Aware Operator Transformer as an Efficient and Accurate Neural Surrogate for PDEs on Arbitrary Domains](https://mlsanity.com/paper/2505.18781/) (2025)
- [Representation Equivalent Neural Operators: a Framework for Alias-free Operator Learning](https://mlsanity.com/paper/2305.19913/) (2023)
- [JAX-FEM: A differentiable GPU-accelerated 3D finite element solver for automatic inverse design and mechanistic data science](https://mlsanity.com/paper/2212.00964/) (2022)
- [Numerical analysis of physics-informed neural networks and related models in physics-informed machine learning](https://mlsanity.com/paper/2402.10926/) (2024)
- [wPINNs: Weak Physics informed neural networks for approximating entropy solutions of hyperbolic conservation laws](https://mlsanity.com/paper/2207.08483/) (2022)
- [RIGNO: A Graph-based framework for robust and accurate operator learning for PDEs on arbitrary domains](https://mlsanity.com/paper/2501.19205/) (2025)
- [Variational Physics Informed Neural Networks: the role of quadratures and test functions](https://mlsanity.com/paper/2109.02035/) (2021)
- [HyPINO: Multi-Physics Neural Operators via HyperPINNs and the Method of Manufactured Solutions](https://mlsanity.com/paper/2509.05117/) (2025)
- [torch-sla: Differentiable Sparse Linear Algebra with Adjoint Solvers and Sparse Tensor Parallelism for PyTorch](https://mlsanity.com/paper/2601.13994/) (2026)

## Built on by (most cited first, 3 of 3)

- [Preconditioned Physics-Informed Neural Operator Training](https://mlsanity.com/paper/2609.36216/) (2026)
- [KernelOnet: An Interpretable Neural Operator Based on Kernel Functions](https://mlsanity.com/paper/2609.35938/) (2026)
- [Transferability of Learned States in Neural PDE Solvers](https://mlsanity.com/paper/2610.10972/) (2026)

## Abstract (arXiv)

We present a unified algorithmic framework for the numerical solution, constrained optimization, and physics-informed learning of PDEs with a variational structure. Our framework is based on a Galerkin discretization of the underlying variational forms, and its high efficiency stems from a novel highly-optimized and GPU-compliant TensorGalerkin framework for linear system assembly (stiffness matrices and load vectors). TensorGalerkin operates by tensorizing element-wise operations within a Python-level Map stage and then performs global reduction with a sparse matrix multiplication that performs message passing on the mesh-induced sparsity graph. The Map and Reduce stages are co-designed inside PyTorch's autograd so that the assembly graph contains O(1) nodes regardless of how the number of elements and local DoFs scale. We validate this O(1)-graph property by deploying TensorGalerkin downstream as i) a highly-efficient numerical PDEs solver, ii) an end-to-end differentiable framework for PDE-constrained optimization, and iii) a physics-informed operator learning algorithm for PDEs. With multiple benchmarks, including 2D and 3D elliptic, parabolic, and hyperbolic PDEs on unstructured meshes, we demonstrate that the proposed framework provides significant computational efficiency and accuracy gains over a variety of baselines in all the targeted downstream applications.