Support for uncertainty propagation?

It would be amazing to be able to specify uncertainties on ICs and have that propagate through to the solutions.
I know this is a very difficult problem in general. Off the top of my head some challenges are:

1. Gaussian uncertainties are hard, but operations on arbitrary (e.g. non-symmetric) distributions often don't even have an analytic form. This could be approximated to 1st order...
2. What's the new API look like?
3. How to support dense solutions?

Thankfully I think at least point 2 has a workable solution. @patrick-kidger, you've written quax to allow for array-ish objects in JAX. My suggestion would be to make a diffeqsolve(y0=) accept quax classes that handle the distribution and its propagation.

Point 1 still remains hard, but there's a still-useful starting point. The simplest "uncertainty" to support isn't even Gaussian but a simple lower and upper bound interval. That would be a good proof of concept but still useful! I know that the same result could be accomplished by doing diffeqsolve twice, but a) the unified API would be a convenience and b) we could hopefully subsequently implement Gaussian and more complex distributions.

To use the opening example from https://docs.kidger.site/diffrax/usage/getting-started/

from diffrax import diffeqsolve, Dopri5, ODETerm, SaveAt, PIDController
from diffrax import Interval

vector_field = lambda t, y, args: -y
term = ODETerm(vector_field)
solver = Dopri5()
saveat = SaveAt(ts=[0., 1., 2., 3.])
stepsize_controller = PIDController(rtol=1e-5, atol=1e-5)

sol = diffeqsolve(term, solver, t0=0, t1=3, dt0=0.1, y0=Interval(0.9, 1.1),  # note the Interval
                             saveat=saveat, stepsize_controller=stepsize_controller)

print(sol.ts)  # DeviceArray([0.   , 1.   , 2.   , 3.    ])
print(sol.ys)  # Interval(...)  # IDK about the internals

As a related note, having quax classes would also enable a nice bundling of arrays of y0 into a MonteCarloMeasurement approximation of an uncertainty distribution:

sol = diffeqsolve(term, solver, t0=0, t1=3, dt0=0.1, y0=MCMeasurement(...),
                             saveat=saveat, stepsize_controller=stepsize_controller)

print(sol.ts)  # DeviceArray([0.   , 1.   , 2.   , 3.    ])
print(sol.ys)  # MCMeasurement(...)
print(sol.ys.mean())  # DeviceArray([1.   , 0.368, 0.135, 0.0498])
print(sol.ys.std())  # DeviceArray([...])

Start by reviewing the diffeqsolve API and the proposed quax integration described in the issue. Compare the interval and MonteCarloMeasurement examples, then determine the API and propagation behavior before implementation. Done would require an agreed design and support for uncertainty-bearing initial conditions through solutions.