Simple PDE solver + Constrained Optimization with JAX by Philip Mocz
A simple example of solving the advection-diffusion equations with JAX and using it in a constrained optimization problem to find initial conditions that yield desired result.
- from
- JAX
- added
- 2026-10-10
- likes
- 0
similar
-
Introduces JAX, M.D., a differentiable physics library which includes simulation environments, interaction potentials, neural networks, and more.
-
Jax in Action manning.com
A hands-on guide to using JAX for deep learning and other mathematically-intensive applications.
-
JAX, M.D. github.com
Accelerated, differential molecular dynamics.
-
JAXopt github.com
Hardware accelerated (GPU/TPU), batchable and differentiable optimizers in JAX.
-
Using JAX to accelerate our research by David Budden and Matteo Hessel deepmind.com
Describes the state of JAX and the JAX ecosystem at DeepMind.
-
JAX-DFT github.com
One-dimensional density functional theory (DFT) in JAX, with implementation of Kohn-Sham equations as regularizer: building prior knowledge into machine-learned physics.
JAX › Tutorials and Blog Posts: “A simple example of solving the advection-diffusion equations with JAX and using it in a constrained optimization problem to find initial conditions that yield desired result.”