This is a follow-up to my Tesseract hackathon honourable mention, tesseract-pinn-inverse-burgers. The original project demonstrated swapping one PINN between JAX and PyTorch. I’ve since rebuilt it around a broader use of Tesseract: packaging swappable, differentiable model components behind typed contracts.
The problem is to recover the viscosity, ν, of the 1D viscous Burgers equation from sparse, noisy measurements. The updated repository solves it in three ways:
- Solver-adjoint: jax.grad differentiates through a spectral Burgers solver Tesseract. Its VJP acts as the PDE adjoint.
- PINN: the same JAX/Optax outer loop trains either a JAX/Equinox or PyTorch PINN through an identical Tesseract contract.
- Amortized posterior: an apply-only flow-matching Tesseract maps one observation to posterior samples over ν and two initial-
condition nuisance parameters.
The project now includes versioned component images, checkpointing, calibration diagnostics, and an interactive Streamlit comparison. I enjoyed extending the project, and I’d appreciate any feedback or comments!
Author: Julian Chan