Possible inconsistency in adjoint state initialisation for time-dependent ODEs

Hello,

I am currently trying to compute the gradient of the ODE:

The ODE is defined indeed only for . I forged the following minimal code to compute the gradient with respect to :

import tensorflow as tf
import tensorflow_probability as tfp


def f(t, y, p):
    """Right-hand-side of the ODE
    """
    return p*tf.sqrt(t)


# initial state and parameter
y0, t0, p = tf.constant(1, dtype=tf.float64), tf.constant(0, dtype=tf.float64), tf.Variable(0.5, dtype=tf.float64)

with tf.GradientTape() as tape:
    ode_solution = tfp.math.ode.DormandPrince().solve(f, t0, y0, solution_times=[8e-3], constants={'p': p})
    gradients = tape.gradient(ode_solution.states, [p])

print(ode_solution.states, gradients)

This example gets stuck in an infinite loop while computing the gradient, I think because the variable derivatives in math/ode/base.py:374 is not a number (in my case, square root of a negative value), and it does not work with DormandPrince(). Instead, the BDF solver seems to work correctly.

I was trying to debug it, and I noticed that it happens when grad_fn invokes adjoint_solver._initialize_solver_internal_state( ... ) on math/ode/base.py:452; under the hood, it is computing an initial state for the adjoint state integration between lines 459 and 464 of math/ode/base.py. In doing so, it assumes a time for the computation of the augmented ODE, and takes the opposite of the time (in math/ode/base.py:368). I guess that this may break time-dependent ODEs.

Is there a specific reason for this choice of the initial time? Or one can set it to something else anyway?

The BDF solver suffers the same issue, but does not get stuck for some reason. It is apparently able to digest such NaN initial value of the derivative.

Reproduce the example with DormandPrince and BDF, then inspect math/ode/base.py around lines 368, 374, and 452-464, especially adjoint_solver._initialize_solver_internal_state. Compare the initial derivative and solver behavior for the time-dependent ODE. Done means the reported initialization inconsistency is understood and the affected solver behavior is addressed or clearly documented.