Improved Integrators

[LarsKue]

This PR:

• Introduces integrator methods for sampling with Flow Matching
  • This is only proof of concept for now, so we are limited to Euler and Runge-Kutta integration.
• Allows dynamic step size selection
  • The initial step size is currently chosen as a constant. This needs to be updated, still.
• Allows choosing the integrator method at sampling time
  • This means we are no longer locked into a single algorithm with the fitted network.
• Temporarily addresses Batched Dynamic Step Size Integration for Differential Equation Models #292 by implementing step size consolidation.
• Fixes Propagate training flag to FlowMatching Integrator #288

The current code for integration is still a little messy due to lack of proper encapsulation. However, we can update this as we go without much impact to the user, if any.

I tested this running Flow Matching on the Two Moons example, and on some toy examples, both in non-compiled and compiled contexts.

[stefanradev93]

That already looks really good! Which integrator do you think should we use as a default or should I run some experiments?

[LarsKue]

@stefanradev93 In my tests, RK45 runs in slightly fewer function evaluations with dynamic step size on. For a static step size, I would still choose Euler, since Flow Matching tends to have straight trajectories. But we can leave RK45 as the default, I think, or even implement more algorithms.

That said, there is still an issue somewhere, presumably with negative step sizes (when start_time > stop_time). I will mark the PR as ready for review again when I have fixed this.