ODE Solvers cannot compute gradients with respect to time

Hi,
I'm using tfp.math.ode.Solver and I would like to compute gradient with respect to time:
d/dt \int_0^t odefunc(T, x) dT = odefunc(t, x).

When odefunc is really simple and contains only constants, this works correctly:

def grad_time_work():
  time = tf.constant(1.)
  x_init = tf.constant(.4)

  w = tf.constant(.5)
  b = tf.constant(.2)

  def odefunc(t, x):
    return tf.nn.sigmoid(w * x + b)

  with tf.GradientTape() as tape:
    tape.watch(time)
    res = tfp.math.ode.DormandPrince().solve(
        odefunc, 0., x_init,
        solution_times=[time],
    )
    x1 = res.states[0]
  grad = tape.gradient(x1, time)
  true_grad = odefunc(time, x1)
  print("grad: ", grad.numpy())
  print("true grad: ", true_grad.numpy())

grad_time_work()

However, when odefunc is an instance of tf.keras.Model, tf.Module, and so forth, gradients with respect to time get None:

def grad_time_error():
  time = tf.constant(1.)
  x_init = tf.constant(.4)

  w = tf.Variable(.5)  # changed only here
  b = tf.Variable(.2)  # tf.constant -> tf.Variable

  def odefunc(t, x):
    return tf.nn.sigmoid(w * x + b)

  with tf.GradientTape() as tape:
    tape.watch(time)
    res = tfp.math.ode.DormandPrince().solve(
        odefunc, 0., x_init,
        solution_times=[time],
    )
    x1 = res.states[0]
  grad = tape.gradient(x1, time)
  true_grad = odefunc(time, x1)
  print("grad: ", grad.numpy())
  print("true: ", true_grad.numpy())

grad_time_error()

===>

AttributeError                            Traceback (most recent call last)
<ipython-input-7-0bd74e9bac74> in <module>()
     21   print("true: ", true_grad.numpy())
     22 
---> 23 grad_time_error()

<ipython-input-7-0bd74e9bac74> in grad_time_error()
     18   grad = tape.gradient(x1, time)
     19   true_grad = odefunc(time, x1)
---> 20   print("grad: ", grad.numpy())
     21   print("true: ", true_grad.numpy())
     22 

AttributeError: 'NoneType' object has no attribute 'numpy'

Google Colab to reproduce this issue here

Is this the desired behavior or are there any workarounds?

Start with the linked Google Colab reproduction and trace tfp.math.ode.DormandPrince().solve when the odefunc closes over tf.Variable values. Compare the constant and variable cases, then verify that the intended gradient behavior or a documented workaround is covered by a focused reproduction.