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!
Github Repo Link
Author: Julian Chan
4 Likes
Hi Julian,
Thanks so much for taking this to the next level and sharing your results! It is great to see such a diverse level of Tesseract usage and demonstration of the power of interoperability. This is very timely as I have just been attending a JAX workshop with today’s topics including PINN training and spectral solvers. I love seeing diffrax shining at pruning implementation complexity down to just the problem-specific details and concretely demonstrating the power of differentiable solvers in attacking inverse problems.
You’ve done a great job with the pseudo-spectral solver which correctly achieves exponential convergence in spatial resolution - the 2/3 filter is the correct approach here. However, you might have noticed you’re taking a huge amount of timesteps for using a 5th order solver. This is because you are not accuracy limited, but stability limited as a result of the stiffness of the diffusion term. Treating this term implicitly with an ImEx solver (e.g. diffrax’s KenCarp solvers) is the canonical solution, and if you track the state in Fourier space you can reduce the number of fft’s from five to three and use a diagonal solver for the implicit term (although you might need to be careful as diffrax’s complex support is limited). Even better you could consider using an exponential solver for the diffusion term with an ETDRK algorithm; unfortunately, this is not offered out of the box in diffrax but Felix Koehler has a really powerful implementation available in exponax that I suggest you check out.
Another cool thing to try is using nonlinear least squares (e.g. optimistix.LevenbergMarquardt) instead of gradient descent. This should work beautifully for the diffrax solver if you switch the adjoint to ForwardMode (as you have one inputsand many residuals). However, for the PINN you might need to play around with increasing collocation points or reducing model size to get this to work as the current setup frames the problem as hugely underdetermined.
Note that for the JAX PINN I don’t think the PINN reconstruction overhead should be material. Yes, this is non-zero, but Equinox tries very hard to minimise this and from what I understand it should be a µs level detail, perhaps there was a performance impact coming from elsewhere? The key unlock performance-wise however is observing that you can reduce that 10/5 separate apply/vjp calls down to one each by concatenating them into a single [obs; col; ic; bc_left; bc_right] batch array (plus folding the per-term loss diagnostics into the same call via has_aux). If you’re interested, we also currently have an experimental tesseract-core PR for caching residuals allowing to avoid some redundant compute. More materially for a 1D diffrax solve specifically, you can probably safely increase the number of checkpoints without running into memory issues.
It would be great to also hear your thoughts on the pros and cons of each method. Where do you think purely differentiable solvers supplant PINN’s and where do PINN’s still provide real value? What is the ideal niche for FMPE posterior approaches?
Thanks so much again for your contribution, inquisitiveness and perseverance. We greatly look forward to hearing from you again in the future!
Hi Jonathan,
Thanks for all the feedback and suggestions, much appreciated!
On your questions:
My understanding is that when the forward operator is known and the parameter to be inferred is low-dimensional (e.g., scalar coefficients), differentiable solvers with discrete adjoints dominate: they restrict optimization to the physical parameter space and enforce the dynamics exactly. In contrast, PINNs introduce a joint state–parameter optimization under soft residual constraints, which significantly degrades conditioning; for canonical cases like Burgers’ viscosity identification, the adjoint approach is strictly superior in accuracy and stability (so in short hard vs soft constraints).
I think PINNs are still well-suited to settings with functional (neural-field) unknowns such as closures or constitutive laws, where the inverse problem is fundamentally infinite-dimensional, or you have some sort of irregular or sparse observation geometries where mesh-based solvers become cumbersome.
Yes, in terms of the FMPE, I’m still fairly new to this. But what I can gather is if I had only have one expensive dataset, I’d use the differentiable solver directly with HMC/NUTS or MAP+Laplace to get a calibrated posterior, rather than training a full amortized model. FMPE becomes somewhat attractive when I expect many repeated inverse problems with the same sensor layout, because I can amortize the simulation and training cost, and each new posterior reduces to a cheap forward pass. One caveat is that FMPE (and flow matching in general) can struggle when the posterior is multimodal or has heavy tails compared to HMC/NUTS, which will find those modes given enough time.
Very timely questions! Some relevant papers you might be interested: Adjoint Method versus Physics-Informed Neural Networks in PDE-Constrained Inverse Problems | alphaXiv , Flow Matching for Scalable Simulation-Based Inference | alphaXiv
Best,
Julian
