JetSCI: A Hybrid JAX-PETSc Framework for Scalable Differentiable Simulation
This work addresses the need for scalable differentiable simulation in scientific machine learning by bridging the gap between automatic differentiation frameworks and high-performance computing libraries.
JetSCI integrates JAX's differentiable programming with PETSc's scalable solvers to enable efficient and accurate differentiable simulation of large-scale finite element problems, outperforming JAX-only implementations in both efficiency and accuracy.
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.