code: def test_riccati_custom_adjoint_solver(
      self, solver, solution_times_fn
  ):
      ode_fn = lambda time, state: (state - time) ** 2 + 1.0
      initial_time = 0.0
      initial_state_value = 0.5
      initial_state = tf.constant(initial_state_value, dtype=tf.float64)
      final_time = 1.0
      solution_times = solution_times_fn(final_time)
      jacobian_fn = lambda time, state: 2.0 * (state - time)
  
      if not isinstance(solution_times, base.ChosenBySolver):
          self.skipTest("b/194468619")
  
      # Instrument the adjoint solver for testing. We have to do this because the
      # API doesn't provide access to the adjoint solver's diagnostics.
      first_step_size = np.float64(1.0)
      solve_count = tf.Variable(0, dtype=tf.int32)
      self.evaluate(solve_count.initializer)
  
      class _InstrumentedSolver(StepSizeHeuristicAdjointSolver):
          def solve(self, **kwargs):
              results = super(_InstrumentedSolver, self).solve(**kwargs)
              with tf.control_dependencies(
                  [solve_count.assign(results.solver_internal_state[1])]
              ):
                  return tf.nest.map_structure(tf.identity, results)
  
      adjoint_solver = _InstrumentedSolver(
          make_solver_fn=lambda step_size: solver(  # pylint: disable=g-long-lambda
              rtol=_RTOL, atol=_ATOL, first_step_size=step_size
          ),
          first_step_size=first_step_size,
      )
  
      solver_instance = solver(
          rtol=_RTOL,
          atol=_ATOL,
          make_adjoint_solver_fn=lambda: adjoint_solver,
      )
  
      def grad_fn(initial_state):
          results = solver_instance.solve(
              ode_fn,
              initial_time,
              initial_state,
              solution_times=solution_times,
              jacobian_fn=jacobian_fn,
          )
          final_state = results.states[-1]
          return final_state
  
      _, grad = tfp_gradient.value_and_gradient(grad_fn, initial_state)
      grad = self.evaluate(grad)
      # There's a race condition if we evaluate solve_count right away. Evaluate
      # it after we're done the computation to produce `grad`.
      solve_count = self.evaluate(solve_count)
      grad_exact = 1.0 / (1.0 - initial_state_value * final_time) ** 2
      self.assertAllClose(grad, grad_exact, rtol=1e-3, atol=1e-3)
      # This indicates that the adaptation carried over to the final solve.
      self.assertGreater(solve_count, 0)
  
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import gradient
  from tensorflow_probability.python.math.ode import base
  from tensorflow_probability.python.math.ode import bdf
  from tensorflow_probability.python.math.ode import dormand_prince
label_desc: Use image node_id 2 for calling the GradientTest obj's underlying member method code with example usage: obj.test_riccati_custom_adjoint_solver(solver, solution_times_fn) and returns: final_state, tf