tensormesh-fem 0.2.0

A fast, differentiable, JIT-free, debugging-friendly finite element library for PyTorch.

Documentation  |  Installation  |  Quickstart  |  Examples  |  Citation

TensorMesh is a finite element method (FEM) library built natively on PyTorch. It is designed to solve partial differential equations (PDEs) with the ergonomics of modern deep learning frameworks — automatic differentiation, GPU acceleration, eager execution — without sacrificing the rigour of classical FEM. Custom weak forms are written in pure Python; the library takes care of tensorized assembly, sparse linear algebra, boundary conditions, and time integration.

Core strengths

• GPU-native & differentiable. Built on PyTorch from the ground up. Moving an entire FEM workflow to the GPU takes a single line of code — every downstream assembly, solve, and gradient inherits the device automatically, with no separate backend or data-marshalling step. Native autograd flows seamlessly through assembly and solve, enabling end-to-end differentiable PDE pipelines.
• High-performance tensorized assembly. A fully tensorized Map-Reduce algorithm powered by TensorGalerkin, which fuses element-wise operations into monolithic GPU kernels, eliminating Python-level loops and delivering order-of-magnitude speedups over CPU-based FEM stacks.
• JIT-free & debugging-friendly. Eager execution with no compilation overhead. Dynamic meshes, adaptive refinement, and interactive workflows just work — no recompilation latency, no opaque traces.
• Comprehensive element & mesh support. Triangular, tetrahedral, pyramid, and prismatic elements with automated mesh generation for common geometries and seamless Gmsh / VTKHDF5 I/O. Multi-field mixed assembly (e.g. LBB-stable Taylor-Hood Stokes) declares each field's order independently of the mesh — quadratic spaces are generated topologically even on linear meshes.
• Flexible Solvers. Powered by torch-sla, our companion library for differentiable sparse linear algebra. Linear, nonlinear, and eigenvalue solvers run across multiple backends on CPU and GPU, with full autograd support, batched solves, and distributed multi-GPU scaling.
• Pythonic API. Custom weak forms in pure Python — no separate DSL, no form compiler. If you can write PyTorch, you can write FEM.

Installation

Requirements: Python ≥ 3.10, PyTorch ≥ 2.0.

pip install tensormesh-fem             # CPU stack (SciPy + PyTorch Krylov)
pip install "tensormesh-fem[gpu]"      # + all PyPI solver backends (cuDSS + PyAMG)

The base install ships the CPU sparse stack via torch-sla — SciPy direct/iterative plus a device-agnostic native-PyTorch Krylov backend that already runs on CUDA/ROCm. If you only want one extra backend, use [cudss] or [pyamg] instead:

pip install "tensormesh-fem[cudss]"    # NVIDIA cuDSS (fastest GPU direct)
pip install "tensormesh-fem[pyamg]"    # PyAMG algebraic multigrid

Two more backends — STRUMPACK (portable direct, CPU/CUDA/ROCm) and NVIDIA AmgX — ship as prebuilt wheels on torch-sla's GitHub Releases, not PyPI. The quotes are needed because [...] is a shell glob character.

git clone https://github.com/camlab-ethz/TensorMesh.git
cd TensorMesh
pip install -e ".[test]"

After installing, sanity-check the install:

python -m tensormesh.verify_install

To see which sparse-solver backends are usable on your machine — and a one-line install hint for any that are not — run:

import torch_sla
torch_sla.show_backends()

Quickstart

Solve $-\Delta u = f$ on the unit square with homogeneous Dirichlet boundary conditions, where $f(x, y) = 2\pi^2\sin(\pi x)\sin(\pi y)$ so the exact solution is $u(x, y) = \sin(\pi x)\sin(\pi y)$:

import math
import torch
from tensormesh import ElementAssembler, NodeAssembler, Mesh, Condenser

# 1. Generate a triangular mesh of the unit square.
mesh = Mesh.gen_rectangle(chara_length=0.05)

# 2. Stiffness weak form:  a(u, v) = ∫ ∇u · ∇v dΩ
class LaplaceAssembler(ElementAssembler):
    def forward(self, gradu, gradv):
        return gradu @ gradv

# 3. Load weak form:  l(v) = ∫ f v dΩ
class SourceAssembler(NodeAssembler):
    def forward(self, v, f):
        return f * v

# 4. Source term, evaluated at every mesh node.
x, y = mesh.points[:, 0], mesh.points[:, 1]
f_vals = 2 * math.pi**2 * torch.sin(math.pi * x) * torch.sin(math.pi * y)

# 5. Assemble the stiffness matrix and load vector.
K = LaplaceAssembler.from_mesh(mesh)()
b = SourceAssembler.from_mesh(mesh)(point_data={"f": f_vals})

# 6. Apply Dirichlet BCs by static condensation, then solve.
condenser = Condenser(mesh.boundary_mask)
K_, b_ = condenser(K, b)
u_ = K_.solve(b_, verbose = True)
u = condenser.recover(u_)

# 7. Compare against the analytical solution.
u_exact = torch.sin(math.pi * x) * torch.sin(math.pi * y)
print(f"L2 error: {(u - u_exact).norm() / u_exact.norm():.3e}")

[torch-sla] solve: n=431, nnz=2859, dtype=float64, device=cpu, symmetric=True, spd=False, backend=scipy, method=lu
L2 error: 3.135e-03

The workflow is Mesh → Assembler → SparseMatrix → Condenser → Solve. Move everything to GPU with a single mesh = mesh.cuda(); enable gradients with mesh.points.requires_grad_(True) and the same script becomes an inverse problem.

See the full walkthrough in the Quickstart.

Examples

A small selection from the example gallery:

3D Poisson — tetrahedral mesh, cut view of the scalar field.	Allen–Cahn phase field — nonlinear time evolution with Newton iteration per step.	Wave equation — explicit central-difference time integration.	Hyperelastic rubber — large-deformation solid mechanics with a Newton solver.
Lid-driven cavity — incompressible Navier–Stokes; velocity field and streamlines.	Magnetostatics — 3D magnetic field around a current-carrying wire (stabilized nodal curl-curl).	Topology optimization — compliance minimization via the Optimality Criteria method.	Physics-informed learning — a network trained to minimize the assembled Galerkin residual.

Feature comparison

Custom Weak Forms (Pythonic) — user-defined bilinear / linear forms directly in Python, without a separate DSL such as UFL. End-to-End Autograd — gradients flow natively through the entire pipeline; FEniCS supports this via the external dolfin-adjoint package. Maturity — reflects project age, ecosystem size, and production deployments.

Architecture

The core workflow: Mesh → Assembler → SparseMatrix → Condenser → Solve.

Documentation

Full documentation, including a user guide, an example gallery, the API reference, and performance benchmarks, lives at docs.tensor-mesh.com.

Key entry points:

• Getting started — installation, quickstart, and an install smoke-test.
• User guide — meshes, weak forms, mixed assembly, boundary conditions, linear solvers, time integration, differentiability, distributed FEM.
• Example gallery — runnable examples from Poisson to Navier–Stokes and topology optimization.
• API reference — module-by-module signatures and docstrings.
• Performance — benchmarks against FEniCS / Firedrake / MFEM / scikit-fem / JAX-FEM / torch-fem.

Community

• Discord — real-time chat, help channels, and showcase.
• GitHub Discussions — announcements, Q&A, ideas & RFCs.
• GitHub Issues — bug reports and feature requests.
• Email to Shizheng Wen at shizheng.wen@sam.math.ethz.ch — collaborations and partnerships.

Contributing

Contributions are welcome — see CONTRIBUTING.md for the development setup, test workflow, documentation build, and PR conventions.

Citation

TensorMesh is the FEM solver component of the TensorGalerkin framework. If you use TensorMesh in your research, please cite the TensorGalerkin paper:

@article{wen2026tensorgalerkin,
  title   = {Learning, Solving and Optimizing PDEs with {TensorGalerkin}:
             an Efficient High-Performance Galerkin Assembly Algorithm},
  author  = {Wen, Shizheng and Chi, Mingyuan and Yu, Tianwei and
             Moseley, Ben and Michelis, Mike Yan and Ren, Pu and
             Sun, Hao and Mishra, Siddhartha},
  journal = {arXiv preprint arXiv:2602.05052},
  year    = {2026}
}

If your work also relies on torch-sla (TensorMesh's solver backend), please additionally cite:

@article{chi2026torchsla,
  title   = {torch-sla: Differentiable Sparse Linear Algebra with Adjoint
             Solvers and Sparse Tensor Parallelism for PyTorch},
  author  = {Chi, Mingyuan and Wen, Shizheng},
  journal = {arXiv preprint arXiv:2601.13994},
  year    = {2026}
}

Acknowledgements