Integrators: Allow passing an integration schedule with fixed time-points

[vpratz]

Some methods (e.g., diffusion models) [1] provide a fixed schedule for integration that takes into account the (approximately known) curvature of the underlying ODE.

Currently, only fixed single step size and (automatic) adaptive step size schedules are possible, but no fixed schedules with varying steps sizes.

@LarsKue How can we introduce them best?
One option woud be the following: If the steps parameter to integrate is iterable and numeric, we interpret it as times for a fixed schedule with [start_time, t_1, t_2, ..., stop_time]. In this setting, we would make the start_time and stop_time parameters optional, and introduce checks to ensure they are passed in the other cases.

What do you think? If you are ok with this design, I would offer to create a PR for this.

[1] https://arxiv.org/abs/2206.00364

[stefanradev93]

@arrjon @niels-leif-bracher are working on this. Perhaps they can chip in.

[vpratz]

@stefanradev93 Thanks for the heads-up. I have preliminary implementation now as well, I think the main question is how we want the interface to look like, i.e., how we want to incorporate into the integrate function.

def integrate_schedule(
    fn: Callable,
    state: dict[str, ArrayLike],
    steps: ArrayLike,
    method: str = "rk45",
    **kwargs,
) -> dict[str, ArrayLike]:

    match method:
        case "euler":
            step_fn = euler_step
        case "rk45":
            step_fn = rk45_step
        case str() as name:
            raise ValueError(f"Unknown integration method name: {name!r}")
        case other:
            raise TypeError(f"Invalid integration method: {other!r}")

    steps = ops.convert_to_tensor(steps)
    step_fn = partial(step_fn, fn, **kwargs, use_adaptive_step_size=False)

    def body(_loop_var, _loop_state):
        
        _state, = _loop_state
        _time = steps[_loop_var]
        step_size = steps[_loop_var + 1] - steps[_loop_var]
        _state, _time, _ = step_fn(_state, _time, step_size)

        return _state,

    state, = keras.ops.fori_loop(0, len(steps) - 1, body, (state,))
    return state

[LarsKue]

Thanks @vpratz this already looks great. I agree that passing the time points as an array is likely the easiest way to define the schedule.

[vpratz]

Thanks for your comment, then I will prepare a PR for this...

[vpratz]

Closed by #404.