JAX/Optax implementation of the Polar Express optimizer (Amsel et al. 2025) applied to PINNs for the first time, plus a KronAdaGrad variant using Polar Express for A **-0.5 preconditioning. Super-exponential convergence over Newton-Schulz.
optimizer muon pinn jax second-order-optimization optax physics-informed-neural-networks polar-decomposition
-
Updated
Sep 20, 2025 - Python