Use jnp.nanmin and jnp.nanmax in jax.experimental.ode.optimal_step_size

[virajpandya]

This simple bug fix prevents jax.experimental.ode.odeint from getting stuck on NaN's if a trial adaptive stepsize led to a NaN for mean_error_ratio.

See the following discussions for more info:
#14612
patrick-kidger/diffrax#223

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[jakevdp]

Looks good! Is there a simple case you could add as a test in tests/ode_test.py?

[jakevdp]

I'm going to close this because it seems to have gone stale. If you'd like to continue with this contribution, feel free to re-open and add a test!