fix for t0==t1 when initializing ys, as opposed to inside the loop

[dkweiss31]

Eliminates code introduced in #494 whereby ys was updated inside of the loop to y0 in the case of t0==t1. Now ys is instead initialized appropriately with y0. This is in response to performance issues observed in #606.

[dkweiss31]

I will investigate these test failures. Patrick can you confirm the adjoint failures are unrelated to this PR? I definitely need to address the event failures though.

[patrick-kidger]

Awesome, thank you for putting this together!

As for the test failures, I think this might be due to these changes being made on top of an old version of Diffrax main. In particular I see the following warning:

/home/runner/work/diffrax/diffrax/diffrax/_integrate.py:1258: DeprecationWarning: shape requires ndarray or scalar arguments, got <class 'jax._src.api.ShapeDtypeStruct'> at position 0. In a future JAX release this will be an error.
1588
    (max_steps,) + jnp.shape(x), jnp.inf, dtype=x.dtype

Which does not reflect the current state or line number of that line:

diffrax/diffrax/_integrate.py

Line 1268 in 2fafbc7

lambda y: jnp.full((out_size,) + y.shape, jnp.inf, dtype=y.dtype), struct

That might explain the failures? (Which I can't reproduce locally.)

[dkweiss31]

Hmm it seems I was wrong: those test failures do directly result from my changes (I do think I based my changes on the most recent version of main, I see commit 2fafbc7 as the most recent change before mine. You might be referring to the below line?)

diffrax/diffrax/_integrate.py

Lines 1257 to 1259 in 36c1d3e

	_make_full = lambda x: jnp.full(
	(max_steps,) + jnp.shape(x), jnp.inf, dtype=x.dtype
	)

In particular, commenting out the below section (which actually makes the replacement with the correct value of y0 if t0==t1), the test_adjoint.py tests pass locally, whereas they do fail locally if I leave this section uncommented. I'm a little baffled by these errors, which don't seem to point to a particular line: maybe its upset that we are asking about t0==t1 in where?

diffrax/diffrax/_integrate.py

Lines 1216 to 1222 in 36c1d3e

	if out_size != 0:
	y0_subsave = jtu.tree_map(
	lambda y: jnp.stack([y] * out_size), subsaveat.fn(t0, y0, args)
	)
	ys = jtu.tree_map(
	lambda _y0, _ys: jnp.where(t0 == t1, _y0, _ys), y0_subsave, ys
	)

[patrick-kidger]

Closing in favour of #618.