taking gradient of ODE solution results in matrix inversion error

I am using tfp.math.ode.BDF to solve a system of ordinary differential equations (ODEs). See my Colaboratory notebook here.

Like the example code in the API documentation, the function ode_fn(t, y, theta) defines the system of ODEs to be solved. I am able to take the gradient of ode_fn wrt theta and integrate the ODEs with tfp.math.ode.BDF.

When I attempt to take the gradient of the ODE solution results wrt theta, however, I get the following error. The code runs without any issues when I replace ode_fn with a simpler set of ODEs. Should the solver settings be adjusted to avoid this error?

InvalidArgumentError                      Traceback (most recent call last)
<ipython-input-9-77ebcb7dd888> in <module>()
----> 1 print(g.gradient(foo, theta0))

5 frames
/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/backprop.py in gradient(self, target, sources, output_gradients, unconnected_gradients)
   1088         output_gradients=output_gradients,
   1089         sources_raw=flat_sources_raw,
-> 1090         unconnected_gradients=unconnected_gradients)
   1091 
   1092     if not self._persistent:

/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/imperative_grad.py in imperative_grad(tape, target, sources, output_gradients, sources_raw, unconnected_gradients)
     75       output_gradients,
     76       sources_raw,
---> 77       compat.as_str(unconnected_gradients.value))

/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py in _backward_function_wrapper(*args)
   1301           break
   1302       return backward._call_flat(  # pylint: disable=protected-access
-> 1303           processed_args, remapped_captures)
   1304 
   1305     return _backward_function_wrapper, recorded_outputs

/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py in _call_flat(self, args, captured_inputs, cancellation_manager)
   1962       # No tape is watching; skip to running the function.
   1963       return self._build_call_outputs(self._inference_function.call(
-> 1964           ctx, args, cancellation_manager=cancellation_manager))
   1965     forward_backward = self._select_forward_and_backward_functions(
   1966         args,

/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py in call(self, ctx, args, cancellation_manager)
    594               inputs=args,
    595               attrs=attrs,
--> 596               ctx=ctx)
    597         else:
    598           outputs = execute.execute_with_cancellation(

/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/execute.py in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     58     ctx.ensure_initialized()
     59     tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
---> 60                                         inputs, attrs, num_outputs)
     61   except core._NotOkStatusException as e:
     62     if name is not None:

InvalidArgumentError:  Input matrix is not invertible.
	 [[{{node gradients/IdentityN_grad/bdfGradients/while/body/_718/gradients/IdentityN_grad/bdfGradients/while/bdf/while/body/_2247/gradients/IdentityN_grad/bdfGradients/while/bdf/while/while/body/_3245/gradients/IdentityN_grad/bdfGradients/while/bdf/while/while/while/body/_5670/gradients/IdentityN_grad/bdfGradients/while/bdf/while/while/while/while/body/_7588/gradients/IdentityN_grad/bdfGradients/while/bdf/while/while/while/while/triangular_solve/MatrixTriangularSolve}}]] [Op:__inference___backward_debug_ode_solver_9192_32890]

Function call stack:
__backward_debug_ode_solver_9192

Start by running the linked Colaboratory notebook and reproducing the gradient call through tfp.math.ode.BDF with the supplied ode_fn and theta. Compare the failing ODE system with the simpler working system and determine whether solver settings prevent the non-invertible matrix error; done means the solution gradient runs successfully without that exception.