Solving differential equations using JAX
JAX provides a set of tools for solving a wide array of problems. Solving differential equations—such as initial value problems described in the Solving simple differential equations numerically recipe—should be well within the capabilities of this library. The diffrax package provides various solvers for differential equations leveraging the power and convenience of JAX.
In the earlier recipe, we solved a relatively simple first-order ODE. In this recipe, we’re going to solve a second-order ODE to illustrate the technique. A second-order ODE is a differential equation that involves both the first and second derivatives of a function. To keep things simple, we’re going to solve a linear second-order ODE of the following form:
Here, 
is a function of 
to be found. In particular, we’re going to solve the following equation:
The initial conditions are 
and...
