GitHub - PredictiveIntelligenceLab/jaxpi
physics-informed-neural-networksjaxscientific-computingpde-solving
Abstraction: JAX library implementing physics-informed neural networks with benchmarks
Key points:
- Comprehensive PINN implementation in JAX supporting single and multi-GPU training; evaluation is single-GPU only
- Integrates architectures from multiple papers including PirateNets (residual adaptive networks), Fourier feature networks, and ActNet (Kolmogorov superposition alternatives)
- Addresses gradient flow pathologies via second-order optimization and neural tangent kernel analysis
- Requires JAX 0.4.36, CUDA 12.4, cuDNN 8.9; default precision set to
highestfor reproducibility - Benchmarks include Allen-Cahn, advection, Stokes flow, Kuramoto-Sivashinsky, Navier-Stokes equations
- Uses Weights & Biases for training metric logging; checkpoints hosted on Google Drive
Connections: Jax · Weights And Biases · Physics Informed Neural Networks · Scientific Machine Learning · Neural Tangent Kernel