ODE solver with documentation minimal example triggers excessive retracing

The example for the ODE solver found in the docs is as follows:

t_init, t0, t1 = 0., 0.5, 1.
y_init = tf.constant([1., 1.], dtype=tf.float64)
A = tf.constant([[-1., -2.], [-3., -4.]], dtype=tf.float64)

def ode_fn(t, y, A):
  return tf.linalg.matvec(A, y)

with tf.GradientTape() as tape:
  tape.watch(A)
  results = tfp.math.ode.BDF().solve(ode_fn, t_init, y_init,
                                     solution_times=[t0, t1],
                                     constants={'A': A})
tape.gradient(results.states, A)  # Fine.

This causes many warnings each with the message:

WARNING:tensorflow:11 out of the last 11 calls to <function pfor..f at 0x7faf2b67db00> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for more details.

I don't think the problem is in ode_fn as there is no shape change there. Perhaps it is in the underlying solver?

Any ideas?

Start by running the minimal ODE example linked in the BDF API documentation and confirm the repeated TensorFlow retracing warnings. Trace the BDF().solve path to determine whether the warning comes from the example or solver, then update the example or solver behavior so the documented run no longer produces excessive retracing warnings.