Integrate a system of ordinary differential equations. Solves the initial value problem for a non-stiff system of first order ODEs: ```

Evaluate polynomial interpolation at the given time point. Args: coefficients: list of Tensor coefficients as created by `interp_fit`.

Take an adaptive Runge-Kutta step to integrate the ODE.

1x1 convolution

3x3 convolution with padding

Automatically links up the gradient from the event time.

Smaller error with slightly more compute.

Updates state based on an event (collision).

Checks that each element is within the error tolerance.

Fit coefficients for 4th order polynomial interpolation. Args: y0: function value at the start of the interval. y1: function valu

Fit an interpolating polynomial to the results of a Runge-Kutta step.

Calculate the optimal size for the next step.

Take an arbitrary Runge-Kutta step and estimate error. Args: func: Function to evaluate like `func(t, y)` to compute the time derivative o

Empirically select a good initial step. The algorithm is described in [1]_. References ---------- .. [1] E. Hairer, S. P. Norsett G.

We ensure all event functions are initially positive, so then we can combine them by taking a min.

Parametric formula for 2d spiral is `r = a + b * theta`. Args: nspiral: number of spirals, i.e. batch dimension ntotal: total number

In-place modifies the adjoint options to choose or wrap the norm function.

Allows training with DataLoaders in a single infinite loop: for i, (x, y) in enumerate(inf_generator(train_loader)):