question about experimental/ode.py

[mariogeiger]

I was looking at experimental/ode.py to change the error criteria and some thing surprised me.

The function error_ratio clearly return the mean of the error ratio: an array of shape ()

https://github.com/google/jax/blob/11127add2bab5bb963a1bacb7eca42559deea6bf/jax/experimental/ode.py#L131-L134

But then when it is used, it is used like if it was an array of ratios: the variable name has an s (error_ratios) and there is an jnp.all as if the intent was to check that all the ratios were smaller than one.

https://github.com/google/jax/blob/11127add2bab5bb963a1bacb7eca42559deea6bf/jax/experimental/ode.py#L198-L204

As another clue that error_ratios was a once list, optimal_step_size contains a jnp.max that has also no effects.

https://github.com/google/jax/blob/11127add2bab5bb963a1bacb7eca42559deea6bf/jax/experimental/ode.py#L136-L139

Question: is this intentional? If so, can I propose a PR to clarify the code?

(Ideally I would like to propose a PR that would allow the user to optionally specify the function error_ratio)

[jakevdp]

It looks like this has returned a scalar since it was first introduced to the JAX package: #1068

I'm not familiar enough with the logic here to know whether it should be an array. Perhaps @duvenaud could comment?

[duvenaud]

I'm pretty sure @mariogeiger is right, and that this is a bug that I introduced when translating this code from TorchDiffEq, because I got confused about shapes at the time.

I actually caught and fixed this little bug a year ago when translating this file to Dex (Dex version here), in which error_ratio returns a Float. Shamefully, I never got around to removing the useless call to np.all in the JAX odeint, because it appears to be harmless. Thank you @mariogeiger for noticing!

[jakevdp]

@mariogeiger are you interested in putting together a PR addressing this?

[mariogeiger]

Sure :-)

[mariogeiger]

#8339