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.

• The paper introduces a hybrid framework that integrates JAX for elementwise autodiff with PETSc for robust, scalable sparse algebra in PDE simulations.
• It employs a matrix-explicit approach with zero-copy device transfers, leveraging JIT compilation and MPI for efficient distributed computation.
• Experimental results show that PETSc-enabled solvers achieve superior convergence and multi-GPU scalability compared to JAX-only implementations.

JetSCI: A Hybrid JAX–PETSc Framework for Scalable Differentiable Simulation

Motivation and Context

Scientific machine learning (SciML) is reshaping computational simulation, intensifying demands for automatic differentiation, surrogate modeling, and scalable PDE solvers within a unified workflow. JAX's automatic differentiation and accelerator support have made it the de facto ML-native environment for local numerical kernels, but JAX is fundamentally constrained in global sparse linear algebra, robust preconditioning, and distributed-memory scalability, which are the hallmarks of mature HPC libraries like PETSc. PETSc provides Krylov solvers, preconditioners, nonlinear solution strategies, and MPI support, but lacks integration with differentiable programming and hardware-accelerated element-level computation. This paper presents JetSCI, a hybrid framework that leverages JAX for differentiable local kernel construction and PETSc for scalable distributed global solution, facilitating multilevel parallelism and efficient memory management.

Framework Design and Software Stack

JetSCI explicitly separates local differentiated computation from global sparse solution. Elementwise residual and Jacobian assembly is performed in JAX, exploiting automatic differentiation, vectorization (via vmap), and JIT compilation. The global Jacobian is assembled in sparse format and transferred—with minimal CPU staging—into PETSc objects for solution. Key interoperability is achieved via CuPy and DLPack to permit zero-copy device-resident array exchange and ctypes for low-level PETSc matrix construction from device pointers. mpi4py mediates distributed-memory execution for scalable parallelism.

(Figure 1)

Figure 1: JetSCI workflow depicting elementwise JAX differentiation, sparse assembly, no-copy handoff, PETSc solver execution, and distributed-memory MPI support.

JetSCI offers two main integration paradigms: matrix-free (via PETSc's Mat SHELL operators evaluated by JAX) and matrix-explicit (explicit sparse assembly in JAX, handoff to PETSc). Empirical results show that the overhead in matrix-free integration dominates for large problem sizes, so JetSCI adopts a matrix-explicit approach to maximize solver robustness and avoid repeated framework boundary crossings.

(Figure 2)

Figure 2: JetSCI software stack, showing the layered integration of JAX, CuPy/DLPack, petsc4py/ctypes, PETSc, and MPI runtime.

Differentiable Assembly and Matrix Handoff

Batch-based FEM assembly enables efficient JIT compilation, as each batch contains elements of identical material, node count, and quadrature. The Jacobian is generated as a dense local block via JAX's autodiff, indexed globally, and then assembled into sparse COO format. Device-resident sparse tuples are converted to CuPy via DLPack, then passed as device pointers for PETSc creation via ctypes, eliminating costly GPU→CPU→GPU transfers.

@jax.vmap def generateJacobian(element): return jax.jacfwd(element.residual_kernel)

Jacobian generation is entirely via JAX's autodiff, sidestepping manual code for complex constitutive laws. The resulting PETSc matrix is ready for MPI-distributed solution and full preconditioning.

Solver Experiments: Performance and Scalability

Extensive experiments benchmark JetSCI against pure-JAX implementations across a series of heterogeneous two-phase micromechanics problems. The global residual and Jacobian system is assembled locally, then solved with iterative Krylov or direct solvers.

Matrix-Vector Product Timings:

Figure 3: Comparison of time to complete 100 matrix-vector products in both PETSc and JAX on a single GPU.

Empirically, PETSc's sparse matrix-vector product throughput outperforms JAX across all problem sizes—a critical consideration for overall solver efficiency.

Iterative Solver Timings and Robustness:

Figure 4: Comparison of time to convergence between PETSc and JAX solvers on a single GPU with Jacobi preconditioning.

JetSCI (PETSc-enabled) solvers outperform JAX-native solvers in convergence time and robustness, particularly as problem size grows. JAX-native solvers diverge at larger scales, while PETSc retains stability.

JAX-only Baselines:

Figure 5: Comparison of time to converge for JAX-only matrix-explicit and matrix-free solvers.

No significant performance difference is observed between JAX-only matrix-explicit and matrix-free paradigms; both are inferior to PETSc-enabled solvers in convergence and scalability.

Direct Solvers:

Figure 6: Convergence time for direct solvers in CuPy and JetSCI.

PETSc's direct solvers (LU/Cholesky), available through JetSCI, remain stable and efficient at problem sizes where CuPy's SuperLU diverges or slows dramatically.

Effect of Preconditioning:

Figure 7: Convergence time for JetSCI solvers with Jacobi and ILU preconditioning.

ILU preconditioning with PETSc provides substantial reductions in convergence time, especially for large, ill-conditioned Jacobians. This capability is only feasible with explicit matrix assembly—a direct advantage of JetSCI's design.

Multi-GPU MPI Scalability:

Figure 8: Convergence time for JetSCI on single GPUs vs multiple GPUs on Notchpeak.

MPI-enabled PETSc objects permit scalable multi-GPU solution. Communication overhead is negligible relative to computational savings at large scales; JetSCI achieves efficient distributed-memory parallelism, unlocking problem sizes unattainable with single-device workflows.

Implications and Future Directions

JetSCI demonstrates that hybrid ML-HPC frameworks can outperform ML-only stack for large-scale, nonlinear, nonsmooth, and heterogeneous PDE simulations. By combining ML-native autodiff for local kernels with MPI-enabled distributed sparse algebra, JetSCI admits scales, robustness, and solver diversity inaccessible in monolithic JAX workflows. The framework avoids unnecessary code duplication and leverages decades of HPC solver development.

Practical implications are immediate for PDE-constrained optimization, differentiable micromechanics, inverse problems, and surrogate modeling. JetSCI's no-copy device data pipeline and structured workflow design offer a template for future SciML systems seeking both productivity and scalability.

Open challenges remain in distributed-memory sparse matrix construction during assembly, further reduction of interop overheads, and richer support for nonlinear and multiphysics workflows. Extending JAX's sharding paradigms, or integrating domain decomposition into JAX-side workflows, could realize even more scalable assembly and inverse modeling.

Conclusion

JetSCI provides a formal integration of ML-native differentiable kernels and HPC-grade sparse algebraic solvers, demonstrating superior efficiency, robustness, and scalability on nonlinear FEM problems with heterogeneous physics. Its design principles—separation of concerns, minimal interop overhead, batchwise differentiated assembly, and explicit sparse handoff—yield strong numerical results and broad practical applicability. Hybrid differentiable-HPC frameworks such as JetSCI will be essential as scientific ML workflows are deployed at ever-increasing scale and complexity.