# JAX-DIPS: Neural bootstrapping of finite discretization methods and application to elliptic problems with discontinuities

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

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Builds on 21 of its 76 references that are in the map.
- First posted on arXiv 25 Oct 2022.

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

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [Neural Operator: Graph Kernel Network for Partial Differential Equations](https://mlsanity.com/paper/2003.03485/) (2020)
- [DGM: A deep learning algorithm for solving partial differential equations](https://mlsanity.com/paper/1708.07469/) (2017)
- [Characterizing possible failure modes in physics-informed neural networks](https://mlsanity.com/paper/2109.01050/) (2021)
- [The Deep Ritz method: A deep learning-based numerical algorithm for solving variational problems](https://mlsanity.com/paper/1710.00211/) (2017)
- [Machine learning accelerated computational fluid dynamics](https://mlsanity.com/paper/2102.01010/) (2021)
- [Model Reduction and Neural Networks for Parametric PDEs](https://mlsanity.com/paper/2005.03180/) (2020)
- [Solver-in-the-Loop: Learning from Differentiable Physics to Interact with Iterative PDE-Solvers](https://mlsanity.com/paper/2007.00016/) (2020)
- [Efficient and Modular Implicit Differentiation](https://mlsanity.com/paper/2105.15183/) (2021)
- [Measuring the Effects of Data Parallelism on Neural Network Training](https://mlsanity.com/paper/1811.03600/) (2018)
- [Learning to Control PDEs with Differentiable Physics](https://mlsanity.com/paper/2001.07457/) (2020)
- [JAX-FLUIDS: A fully-differentiable high-order computational fluid dynamics solver for compressible two-phase flows](https://mlsanity.com/paper/2203.13760/) (2022)
- [Solving inverse-PDE problems with physics-aware neural networks](https://mlsanity.com/paper/2001.03608/) (2020)
- [Extracting Interpretable Physical Parameters from Spatiotemporal Systems using Unsupervised Learning](https://mlsanity.com/paper/1907.06011/) (2019)
- [Data driven approximation of parametrized PDEs by Reduced Basis and Neural Networks](https://mlsanity.com/paper/1904.01514/) (2019)
- [A Discontinuity Capturing Shallow Neural Network for Elliptic Interface Problems](https://mlsanity.com/paper/2106.05587/) (2021)
- [Statistical Properties of the log-cosh Loss Function Used in Machine Learning](https://mlsanity.com/paper/2208.04564/) (2022)
- [Neural network augmented inverse problems for PDEs](https://mlsanity.com/paper/1712.09685/) (2017)
- [A three layer neural network can represent any multivariate function](https://mlsanity.com/paper/2012.03016/) (2020)
- [A non-parametric shape optimization approach for solving inverse problems in directed self-assembly of block copolymers](https://mlsanity.com/paper/2112.09615/) (2021)

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

- [RandONet: Shallow-Networks with Random Projections for learning linear and nonlinear operators](https://mlsanity.com/paper/2406.05470/) (2024)
- [GoRINNs: Godunov-Riemann Informed Neural Networks for Learning Hyperbolic Conservation Laws](https://mlsanity.com/paper/2410.22193/) (2024)
- [JAX-MPM: A Learning-Augmented Differentiable Meshfree Framework for GPU-Accelerated Lagrangian Simulation and Geophysical Inverse Modeling](https://mlsanity.com/paper/2507.04192/) (2025)
- [Neural-Integrated Meshfree (NIM) Method: A differentiable programming-based hybrid solver for computational mechanics](https://mlsanity.com/paper/2311.12915/) (2023)
- [A Multiple Transferable Neural Network Method with Domain Decomposition for Elliptic Interface Problems](https://mlsanity.com/paper/2502.19893/) (2025)
- [Zero Coordinate Shift: Whetted Automatic Differentiation for Physics-informed Operator Learning](https://mlsanity.com/paper/2311.00860/) (2023)
- [Invariant deep neural networks under the finite group for solving partial differential equations](https://mlsanity.com/paper/2407.20560/) (2024)

## Abstract (arXiv)

We present a scalable strategy for development of mesh-free hybrid neuro-symbolic partial differential equation solvers based on existing mesh-based numerical discretization methods. Particularly, this strategy can be used to efficiently train neural network surrogate models of partial differential equations by (i) leveraging the accuracy and convergence properties of advanced numerical methods, solvers, and preconditioners, as well as (ii) better scalability to higher order PDEs by strictly limiting optimization to first order automatic differentiation. The presented neural bootstrapping method (hereby dubbed NBM) is based on evaluation of the finite discretization residuals of the PDE system obtained on implicit Cartesian cells centered on a set of random collocation points with respect to trainable parameters of the neural network. Importantly, the conservation laws and symmetries present in the bootstrapped finite discretization equations inform the neural network about solution regularities within local neighborhoods of training points. We apply NBM to the important class of elliptic problems with jump conditions across irregular interfaces in three spatial dimensions. We show the method is convergent such that model accuracy improves by increasing number of collocation points in the domain and predonditioning the residuals. We show NBM is competitive in terms of memory and training speed with other PINN-type frameworks. The algorithms presented here are implemented using textttJAX in a software package named textttJAX-DIPS (https://github.com/JAX-DIPS/JAX-DIPS), standing for differentiable interfacial PDE solver. We open sourced textttJAX-DIPS to facilitate research into use of differentiable algorithms for developing hybrid PDE solvers.