Running diffeqsolve for a fixed number of *accepted* steps

[BenjaminDAnjou]

Is it possible to run diffeqsolve with an adaptive stepsize controller for a fixed number of accepted steps? The solution diffeqsolve(..., max_steps=n, throw=False, t1=jnp.inf) does not distinguish between rejected and accepted steps.

The solution diffeqsolve(..., max_steps=n, throw=False, t1=jnp.inf) can be put in a subloop that checks sol.stats["num_rejected_steps"]. I tried this, but I think this may require explicitly calling adapt_step_size() outside of diffeqsolve. If I pass only the controller_state to diffeqsolve, I get stuck in an infinite loop whenever a rejection occurs.

[patrick-kidger]

Haha! Good question. I think putting diffeqsolve in a loop is reasonable way to do this. It should be possible to chain multiple solves together as if they were one, by forwarding these three variables:

diffrax/diffrax/_integrate.py

Lines 986 to 994 in 5ea5fcb

	- `solver_state`: Some initial state for the solver. Generally obtained by
	`SaveAt(solver_state=True)` from a previous solve.
	- `controller_state`: Some initial state for the step size controller. Generally
	obtained by `SaveAt(controller_state=True)` from a previous solve.
	- `made_jump`: Whether a jump has just been made at `t0`. Used to update
	`solver_state` (if passed). Generally obtained by `SaveAt(made_jump=True)`
	from a previous solve.

You shouldn't need to manually call adapt_step_size or similar, just keep calling diffeqsolve until you've used up your budget of steps.