underPINN: A Modular and Scalable JAX-Based Framework for Physics-Informed Machine Learning

Physics-informed neural network (PINN) libraries have proliferated across PyTorch, JAX, and Julia, and cross-framework speed comparisons between them are increasingly common. We show that these comparisons are surprisingly fragile. Framework defaults that never appear in a problem's configuration file — weight normalization applied to every hidden layer, renormalization of the network's time input, dimension-restricted differential operators silently violate the “same architecture, same physics” assumption such comparisons rest on. Separately, we find that a single-run PINN ablation can reverse its conclusion under an identical random seed and identical code, as GPU floating-point non-determinism compounds through adaptive-collocation training dynamics. We present underPINN, an open-source JAX-native PINN framework, and use it as a controlled instrument for rigorously matched comparison. underPINN compiles each problem's optimization step — forward evaluation, automatic differentiation, and the optimizer update — once with jax.jit and reuses that executable every epoch; an optional jax.lax.scan-fused path yields no measurable further speedup once the step is compiled, which we report. Against eager-mode PyTorch on three matched reimplementations the compiled path is $1.2$–$22\times$ faster, with the speedup tracking the ratio of per-epoch dispatch overhead to per-step compute rather than problem complexity. Against NVIDIA PhysicsNeMo Sym (eight matched problems) and jinns (six), underPINN is faster on every problem but only after the confounds above are identified and corrected, which we document as guidance for the field rather than a criticism of any one library. Through one configuration interface the framework also exposes Fourier-feature embeddings, gated MLPs, FBPINN domain decomposition, residual-adaptive collocation, a deterministic DEIM-inspired collocation selector (QR-DEIM-R), and a Levenberg–Marquardt-damped Gauss–Newton step. We test these across multiple seeds rather than assume they help, and report what the tests return. Residual-adaptive collocation improves accuracy only where a localized, uncapped feature exists to exploit, not on smooth or dissipation-capped problems. QR-DEIM-R does not reproduce the gains published for QR-DEIM-based collocation on any of the six problems we run it on, and tuned Adam matches or beats the Gauss–Newton step on both accuracy and wall-clock. Our nine problems span ODEs, elliptic and parabolic PDEs, 3-D pipe flow, and 1-/2-D compressible shock flows, and we report the failures alongside the successes including a 1-D blast wave whose relative $L^2$ error plateaus near $33\%$. The compressible cases are solved data-free, from the governing equations and boundary conditions alone with no reference solution supplied during training: among them a Mach 3 viscous shock/boundary-layer-interaction ramp, and Toro Test 3, whose $10^5$ initial pressure ratio sits four orders of magnitude above the Sod- and Lax-type Riemann problems usual in the PINN shock literature.