JetSCI: A Hybrid JAX-PETSc Framework for Scalable Differentiable Simulation
Abstract
The rapid rise of scientific machine learning (SciML) has expanded the role of differentiable modeling, surrogate modeling, and data-driven constitutive laws in large-scale simulation. The JAX framework provides an attractive environment for these workflows through automatically differentiable programs, vectorization, GPU acceleration, and while enabling seamless learning of surrogate models. However, large-scale simulation still relies on mature HPC infrastructure. Libraries, such as PETSc, provide scalable MPI-based parallelism, robust linear and nonlinear solvers, and advanced preconditioning capabilities that remain difficult to reproduce in JAX-only workflows. We present JetSCI, a hybrid JAX-PETSc framework that unifies these complementary strengths. JetSCI uses JAX for GPU-parallel differentiable discretizations and PETSc for robust, scalable solution of the resulting systems on distributed-memory architectures, exposing multilevel parallelism through GPU acceleration within nodes and MPI parallelism across nodes. For finite element discretizations of heterogeneous micromechanics problems, JetSCI outperforms JAX-only implementations in efficiency and accuracy.
Cite
@article{arxiv.2604.22087,
title = {JetSCI: A Hybrid JAX-PETSc Framework for Scalable Differentiable Simulation},
author = {Alberto Cattaneo and M Keith Ballard and Robert M. Kirby and Varun Shankar},
journal= {arXiv preprint arXiv:2604.22087},
year = {2026}
}