class_docstr: Base class for an ODE solver.
class_docstr_tok: ["Base","class","for","an","ODE","solver","."]
docstr: Solves an initial value problem.
  
  An initial value problem consists of a system of ODEs and an initial
  condition:
  
  ```none
  dy/dt(t) = ode_fn(t, y(t), **constants)
  y(initial_time) = initial_state
  ```
  
  Here, `t` (also called time) is a scalar float `Tensor` and `y(t)` (also
  called the state at time `t`) is an N-D float or complex `Tensor`.
  `constants` is are values that are constant with respect to time. Passing
  the constants here rather than just closing over them in `ode_fn` is only
  necessary if you want gradients with respect to these values.
  
  ### Example
  
  The ODE `dy/dt(t) = dot(A, y(t))` is solved below.
  
  ```python
  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):
    return tf.linalg.matvec(A, y)
  
  results = tfp.math.ode.BDF().solve(ode_fn, t_init, y_init,
                                     solution_times=[t0, t1])
  y0 = results.states[0]  # == dot(matrix_exp(A * t0), y_init)
  y1 = results.states[1]  # == dot(matrix_exp(A * t1), y_init)
  ```
  
  If the exact solution times are not important, it can be much
  more efficient to let the solver choose them using
  `solution_times=tfp.math.ode.ChosenBySolver(final_time=1.)`.
  This yields the state at various times between `t_init` and `final_time`,
  in which case `results.states[i]` is the state at time `results.times[i]`.
  
  #### Gradients
  
  The gradients are computed using the adjoint sensitivity method described in
  [Chen et al. (2018)][1].
  
  ```python
  grad = tf.gradients(y1, y0) # == dot(e, J)
  # J is the Jacobian of y1 with respect to y0. In this case, J = exp(A * t1).
  # e = [1, ..., 1] is the row vector of ones.
  ```
  
  This is not capable of computing gradients with respect to values closed
  over by `ode_fn`, e.g., in the example above:
  
  ```python
  def ode_fn(t, y):
    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])
  tape.gradient(results.states, A)  # Undefined!
  ```
  
  There are two options to get the gradients flowing through these values:
  
  1. Use `tf.Variable` for these values.
  2. Pass the values in explicitly using the `constants` argument:
  
  ```python
  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.
  ```
  
  By default, this uses the same solver for the augmented ODE. This can be
  controlled via `make_adjoint_solver_fn`.
  
  #### References
  
  [1]: Chen, Tian Qi, et al. "Neural ordinary differential equations."
       Advances in Neural Information Processing Systems. 2018.
  
  Args:
    ode_fn: Function of the form `ode_fn(t, y, **constants)`. The input `t` is
      a scalar float `Tensor`. The input `y` and output are both `Tensor`s
      with the same shape and `dtype` as `initial_state`. `constants` is are
      values that are constant with respect to time. Passing the constants
      here rather than just closing over them in `ode_fn` is only necessary if
      you want gradients with respect to these values.
    initial_time: Scalar float `Tensor` specifying the initial time.
    initial_state: N-D float or complex `Tensor` specifying the initial state.
      The `dtype` of `initial_state` must be complex for problems with
      complex-valued states (even if the initial state is real).
    solution_times: 1-D float `Tensor` specifying a list of times. The solver
      stores the computed state at each of these times in the returned
      `Results` object. Must satisfy `initial_time <= solution_times[0]` and
      `solution_times[i] < solution_times[i+1]`. Alternatively, the user can
      pass `tfp.math.ode.ChosenBySolver(final_time)` where `final_time` is a
      scalar float `Tensor` satisfying `initial_time < final_time`. Doing so
      requests that the solver automatically choose suitable times up to and
      including `final_time` at which to store the computed state.
    jacobian_fn: Optional function of the form `jacobian_fn(t, y)`. The input
      `t` is a scalar float `Tensor`. The input `y` has the same shape and
      `dtype` as `initial_state`. The output is a 2N-D `Tensor` whose shape is
      `initial_state.shape + initial_state.shape` and whose `dtype` is the
      same as `initial_state`. In particular, the `(i1, ..., iN, j1, ...,
      jN)`-th entry of `jacobian_fn(t, y)` is the derivative of the `(i1, ...,
      iN)`-th entry of `ode_fn(t, y)` with respect to the `(j1, ..., jN)`-th
      entry of `y`. If this argument is left unspecified, the solver
      automatically computes the Jacobian if and when it is needed.
      Default value: `None`.
    jacobian_sparsity: Optional 2N-D boolean `Tensor` whose shape is
      `initial_state.shape + initial_state.shape` specifying the sparsity
      pattern of the Jacobian. This argument is ignored if `jacobian_fn` is
      specified.
      Default value: `None`.
    batch_ndims: Optional nonnegative integer. When specified, the first
      `batch_ndims` dimensions of `initial_state` are batch dimensions.
      Default value: `None`.
    previous_solver_internal_state: Optional solver-specific argument used to
      warm-start this invocation of `solve`.
      Default value: `None`.
    constants: Optional dictionary with string keys and values being (possibly
      nested) float `Tensor`s. These represent values that are constant with
      respect to time. Specifying these here allows the adjoint sentitivity
      method to compute gradients of the results with respect to these values.
  
  Returns:
    Object of type `Results`.
docstr_tok: ["Solves","an","initial","value","problem",".","An","initial","value","problem","consists","of","a","system","of","ODEs","and","an","initial","condition",":","``","`","none","dy\/dt","(","t",")","=","ode_fn","(","t",",","y","(","t",")",",","*","*","constants",")","y","(","initial_time",")","=","initial_state","``","`","Here",",","`","t","`","(","also","called","time",")","is","a","scalar","float","`","Tensor","`","and","`","y","(","t",")","`","(","also","called","the","state","at","time","`","t","`",")","is","an","N-D","float","or","complex","`","Tensor","`",".","`","constants","`","is","are","values","that","are","constant","with","respect","to","time",".","Passing","the","constants","here","rather","than","just","closing","over","them","in","`","ode_fn","`","is","only","necessary","if","you","want","gradients","with","respect","to","these","values",".","#","#","#","Example","The","ODE","`","dy\/dt","(","t",")","=","dot","(","A",",","y","(","t",")",")","`","is","solved","below",".","``","`","python","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",")",":","return","tf.linalg.matvec","(","A",",","y",")","results","=","tfp.math.ode.BDF","(",")",".solve","(","ode_fn",",","t_init",",","y_init",",","solution_times=","[","t0",",","t1","]",")","y0","=","results.states","[","0","]","#","==","dot","(","matrix_exp","(","A","*","t0",")",",","y_init",")","y1","=","results.states","[","1","]","#","==","dot","(","matrix_exp","(","A","*","t1",")",",","y_init",")","``","`","If","the","exact","solution","times","are","not","important",",","it","can","be","much","more","efficient","to","let","the","solver","choose","them","using","`","solution_times=tfp.math.ode.ChosenBySolver","(","final_time=1.",")","`",".","This","yields","the","state","at","various","times","between","`","t_init","`","and","`","final_time","`",",","in","which","case","`","results.states","[","i","]","`","is","the","state","at","time","`","results.times","[","i","]","`",".","#","#","#","#","Gradients","The","gradients","are","computed","using","the","adjoint","sensitivity","method","described","in","[","Chen","et","al",".","(","2018",")","]","[","1","]",".","``","`","python","grad","=","tf.gradients","(","y1",",","y0",")","#","==","dot","(","e",",","J",")","#","J","is","the","Jacobian","of","y1","with","respect","to","y0",".","In","this","case",",","J","=","exp","(","A","*","t1",")",".","#","e","=","[","1",",","...",",","1","]","is","the","row","vector","of","ones",".","``","`","This","is","not","capable","of","computing","gradients","with","respect","to","values","closed","over","by","`","ode_fn","`",",","e.g.",",","in","the","example","above",":","``","`","python","def","ode_fn","(","t",",","y",")",":","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","]",")","tape.gradient","(","results.states",",","A",")","#","Undefined","!","``","`","There","are","two","options","to","get","the","gradients","flowing","through","these","values",":","1",".","Use","`","tf.Variable","`","for","these","values",".","2",".","Pass","the","values","in","explicitly","using","the","`","constants","`","argument",":","``","`","python","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",".","``","`","By","default",",","this","uses","the","same","solver","for","the","augmented","ODE",".","This","can","be","controlled","via","`","make_adjoint_solver_fn","`",".","#","#","#","#","References","[","1","]",":","Chen",",","Tian","Qi",",","et","al",".","``","Neural","ordinary","differential","equations",".","''","Advances","in","Neural","Information","Processing","Systems",".","2018",".","Args",":","ode_fn",":","Function","of","the","form","`","ode_fn","(","t",",","y",",","*","*","constants",")","`",".","The","input","`","t","`","is","a","scalar","float","`","Tensor","`",".","The","input","`","y","`","and","output","are","both","`","Tensor","`","s","with","the","same","shape","and","`","dtype","`","as","`","initial_state","`",".","`","constants","`","is","are","values","that","are","constant","with","respect","to","time",".","Passing","the","constants","here","rather","than","just","closing","over","them","in","`","ode_fn","`","is","only","necessary","if","you","want","gradients","with","respect","to","these","values",".","initial_time",":","Scalar","float","`","Tensor","`","specifying","the","initial","time",".","initial_state",":","N-D","float","or","complex","`","Tensor","`","specifying","the","initial","state",".","The","`","dtype","`","of","`","initial_state","`","must","be","complex","for","problems","with","complex-valued","states","(","even","if","the","initial","state","is","real",")",".","solution_times",":","1-D","float","`","Tensor","`","specifying","a","list","of","times",".","The","solver","stores","the","computed","state","at","each","of","these","times","in","the","returned","`","Results","`","object",".","Must","satisfy","`","initial_time","<","=","solution_times","[","0","]","`","and","`","solution_times","[","i","]","<","solution_times","[","i+1","]","`",".","Alternatively",",","the","user","can","pass","`","tfp.math.ode.ChosenBySolver","(","final_time",")","`","where","`","final_time","`","is","a","scalar","float","`","Tensor","`","satisfying","`","initial_time","<","final_time","`",".","Doing","so","requests","that","the","solver","automatically","choose","suitable","times","up","to","and","including","`","final_time","`","at","which","to","store","the","computed","state",".","jacobian_fn",":","Optional","function","of","the","form","`","jacobian_fn","(","t",",","y",")","`",".","The","input","`","t","`","is","a","scalar","float","`","Tensor","`",".","The","input","`","y","`","has","the","same","shape","and","`","dtype","`","as","`","initial_state","`",".","The","output","is","a","2N-D","`","Tensor","`","whose","shape","is","`","initial_state.shape","+","initial_state.shape","`","and","whose","`","dtype","`","is","the","same","as","`","initial_state","`",".","In","particular",",","the","`","(","i1",",","...",",","iN",",","j1",",","...",",","jN",")","`","-th","entry","of","`","jacobian_fn","(","t",",","y",")","`","is","the","derivative","of","the","`","(","i1",",","...",",","iN",")","`","-th","entry","of","`","ode_fn","(","t",",","y",")","`","with","respect","to","the","`","(","j1",",","...",",","jN",")","`","-th","entry","of","`","y","`",".","If","this","argument","is","left","unspecified",",","the","solver","automatically","computes","the","Jacobian","if","and","when","it","is","needed",".","Default","value",":","`","None","`",".","jacobian_sparsity",":","Optional","2N-D","boolean","`","Tensor","`","whose","shape","is","`","initial_state.shape","+","initial_state.shape","`","specifying","the","sparsity","pattern","of","the","Jacobian",".","This","argument","is","ignored","if","`","jacobian_fn","`","is","specified",".","Default","value",":","`","None","`",".","batch_ndims",":","Optional","nonnegative","integer",".","When","specified",",","the","first","`","batch_ndims","`","dimensions","of","`","initial_state","`","are","batch","dimensions",".","Default","value",":","`","None","`",".","previous_solver_internal_state",":","Optional","solver-specific","argument","used","to","warm-start","this","invocation","of","`","solve","`",".","Default","value",":","`","None","`",".","constants",":","Optional","dictionary","with","string","keys","and","values","being","(","possibly","nested",")","float","`","Tensor","`","s",".","These","represent","values","that","are","constant","with","respect","to","time",".","Specifying","these","here","allows","the","adjoint","sentitivity","method","to","compute","gradients","of","the","results","with","respect","to","these","values",".","Returns",":","Object","of","type","`","Results","`","."]
code: def solve(
      self,
      ode_fn,
      initial_time,
      initial_state,
      solution_times,
      jacobian_fn=None,
      jacobian_sparsity=None,
      batch_ndims=None,
      previous_solver_internal_state=None,
      constants=None,
  ):
      """Solves an initial value problem.
  
      An initial value problem consists of a system of ODEs and an initial
      condition:
  
      ```none
      dy/dt(t) = ode_fn(t, y(t), **constants)
      y(initial_time) = initial_state
      ```
  
      Here, `t` (also called time) is a scalar float `Tensor` and `y(t)` (also
      called the state at time `t`) is an N-D float or complex `Tensor`.
      `constants` is are values that are constant with respect to time. Passing
      the constants here rather than just closing over them in `ode_fn` is only
      necessary if you want gradients with respect to these values.
  
      ### Example
  
      The ODE `dy/dt(t) = dot(A, y(t))` is solved below.
  
      ```python
      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):
        return tf.linalg.matvec(A, y)
  
      results = tfp.math.ode.BDF().solve(ode_fn, t_init, y_init,
                                         solution_times=[t0, t1])
      y0 = results.states[0]  # == dot(matrix_exp(A * t0), y_init)
      y1 = results.states[1]  # == dot(matrix_exp(A * t1), y_init)
      ```
  
      If the exact solution times are not important, it can be much
      more efficient to let the solver choose them using
      `solution_times=tfp.math.ode.ChosenBySolver(final_time=1.)`.
      This yields the state at various times between `t_init` and `final_time`,
      in which case `results.states[i]` is the state at time `results.times[i]`.
  
      #### Gradients
  
      The gradients are computed using the adjoint sensitivity method described in
      [Chen et al. (2018)][1].
  
      ```python
      grad = tf.gradients(y1, y0) # == dot(e, J)
      # J is the Jacobian of y1 with respect to y0. In this case, J = exp(A * t1).
      # e = [1, ..., 1] is the row vector of ones.
      ```
  
      This is not capable of computing gradients with respect to values closed
      over by `ode_fn`, e.g., in the example above:
  
      ```python
      def ode_fn(t, y):
        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])
      tape.gradient(results.states, A)  # Undefined!
      ```
  
      There are two options to get the gradients flowing through these values:
  
      1. Use `tf.Variable` for these values.
      2. Pass the values in explicitly using the `constants` argument:
  
      ```python
      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.
      ```
  
      By default, this uses the same solver for the augmented ODE. This can be
      controlled via `make_adjoint_solver_fn`.
  
      #### References
  
      [1]: Chen, Tian Qi, et al. "Neural ordinary differential equations."
           Advances in Neural Information Processing Systems. 2018.
  
      Args:
        ode_fn: Function of the form `ode_fn(t, y, **constants)`. The input `t` is
          a scalar float `Tensor`. The input `y` and output are both `Tensor`s
          with the same shape and `dtype` as `initial_state`. `constants` is are
          values that are constant with respect to time. Passing the constants
          here rather than just closing over them in `ode_fn` is only necessary if
          you want gradients with respect to these values.
        initial_time: Scalar float `Tensor` specifying the initial time.
        initial_state: N-D float or complex `Tensor` specifying the initial state.
          The `dtype` of `initial_state` must be complex for problems with
          complex-valued states (even if the initial state is real).
        solution_times: 1-D float `Tensor` specifying a list of times. The solver
          stores the computed state at each of these times in the returned
          `Results` object. Must satisfy `initial_time <= solution_times[0]` and
          `solution_times[i] < solution_times[i+1]`. Alternatively, the user can
          pass `tfp.math.ode.ChosenBySolver(final_time)` where `final_time` is a
          scalar float `Tensor` satisfying `initial_time < final_time`. Doing so
          requests that the solver automatically choose suitable times up to and
          including `final_time` at which to store the computed state.
        jacobian_fn: Optional function of the form `jacobian_fn(t, y)`. The input
          `t` is a scalar float `Tensor`. The input `y` has the same shape and
          `dtype` as `initial_state`. The output is a 2N-D `Tensor` whose shape is
          `initial_state.shape + initial_state.shape` and whose `dtype` is the
          same as `initial_state`. In particular, the `(i1, ..., iN, j1, ...,
          jN)`-th entry of `jacobian_fn(t, y)` is the derivative of the `(i1, ...,
          iN)`-th entry of `ode_fn(t, y)` with respect to the `(j1, ..., jN)`-th
          entry of `y`. If this argument is left unspecified, the solver
          automatically computes the Jacobian if and when it is needed.
          Default value: `None`.
        jacobian_sparsity: Optional 2N-D boolean `Tensor` whose shape is
          `initial_state.shape + initial_state.shape` specifying the sparsity
          pattern of the Jacobian. This argument is ignored if `jacobian_fn` is
          specified.
          Default value: `None`.
        batch_ndims: Optional nonnegative integer. When specified, the first
          `batch_ndims` dimensions of `initial_state` are batch dimensions.
          Default value: `None`.
        previous_solver_internal_state: Optional solver-specific argument used to
          warm-start this invocation of `solve`.
          Default value: `None`.
        constants: Optional dictionary with string keys and values being (possibly
          nested) float `Tensor`s. These represent values that are constant with
          respect to time. Specifying these here allows the adjoint sentitivity
          method to compute gradients of the results with respect to these values.
  
      Returns:
        Object of type `Results`.
      """
      if constants is None:
          constants = {}
      input_state_structure = initial_state
      initial_state, constants = tf.nest.map_structure(
          tf.convert_to_tensor, (initial_state, constants)
      )
  
      def vjp_fwd(initial_state, constants):
          results = self._solve(
              ode_fn=functools.partial(ode_fn, **constants),
              initial_time=initial_time,
              initial_state=initial_state,
              solution_times=solution_times,
              jacobian_fn=jacobian_fn,
              jacobian_sparsity=jacobian_sparsity,
              batch_ndims=batch_ndims,
              previous_solver_internal_state=previous_solver_internal_state,
          )
          results = Results(
              times=tf.stop_gradient(results.times),
              states=results.states,
              diagnostics=util.stop_gradient_of_real_or_complex_entries(
                  results.diagnostics
              ),
              solver_internal_state=util.stop_gradient_of_real_or_complex_entries(
                  results.solver_internal_state
              ),
          )
          return results, (results, constants)
  
      def vjp_bwd(results_constants, dresults, variables=()):
          """Adjoint sensitivity method to compute gradients."""
          results, constants = results_constants
          adjoint_solver = self._make_adjoint_solver_fn()
          dstates = dresults.states
          # TODO(b/138304303): Support complex types.
          with tf.name_scope("{}Gradients".format(self._name)):
              get_dtype = lambda x: x.dtype
  
              def error_if_complex(dtype):
                  if dtype_util.is_complex(dtype):
                      raise NotImplementedError(
                          "The adjoint sensitivity method does "
                          "not support complex dtypes."
                      )
  
              state_dtypes = tf.nest.map_structure(
                  get_dtype, initial_state
              )
              tf.nest.map_structure(error_if_complex, state_dtypes)
              common_state_dtype = dtype_util.common_dtype(
                  initial_state
              )
              real_dtype = dtype_util.real_dtype(common_state_dtype)
  
              # We add initial_time to ensure that we know where to stop.
              result_times = tf.concat(
                  [[tf.cast(initial_time, real_dtype)], results.times],
                  0,
              )
              num_result_times = tf.size(result_times)
  
              # First two components correspond to reverse and adjoint states.
              # the last two component is adjoint state for variables and constants.
              terminal_augmented_state = tuple(
                  [
                      rk_util.nest_constant(initial_state, 0.0),
                      rk_util.nest_constant(initial_state, 0.0),
                      tuple(
                          rk_util.nest_constant(variable, 0.0)
                          for variable in variables
                      ),
                      rk_util.nest_constant(constants, 0.0),
                  ]
              )
  
              # The XLA compiler does not compile code which slices/indexes using
              # integer `Tensor`s. `TensorArray`s are used to get around this.
              result_time_array = tf.TensorArray(
                  results.times.dtype,
                  clear_after_read=False,
                  size=num_result_times,
                  element_shape=[],
              ).unstack(result_times)
  
              # TensorArray shape should not include time dimension, hence shape[1:]
              result_state_arrays = [
                  tf.TensorArray(  # pylint: disable=g-complex-comprehension
                      dtype=component.dtype,
                      size=num_result_times - 1,
                      clear_after_read=False,
                      element_shape=component.shape[1:],
                  ).unstack(
                      component
                  )
                  for component in tf.nest.flatten(results.states)
              ]
              result_state_arrays = tf.nest.pack_sequence_as(
                  results.states, result_state_arrays
              )
              dresult_state_arrays = [
                  tf.TensorArray(  # pylint: disable=g-complex-comprehension
                      dtype=component.dtype,
                      size=num_result_times - 1,
                      clear_after_read=False,
                      element_shape=component.shape[1:],
                  ).unstack(
                      component
                  )
                  for component in tf.nest.flatten(dstates)
              ]
              dresult_state_arrays = tf.nest.pack_sequence_as(
                  results.states, dresult_state_arrays
              )
  
              def augmented_ode_fn(backward_time, augmented_state):
                  """Dynamics function for the augmented system.
  
                  Describes a differential equation that evolves the augmented state
                  backwards in time to compute gradients using the adjoint method.
                  Augmented state consists of 4 components `(state, adjoint_state,
                  vars, constants)` all evaluated at time `backward_time`:
  
                  state: represents the solution of user provided `ode_fn`. The
                    structure coincides with the `initial_state`.
                  adjoint_state: represents the solution of the adjoint sensitivity
                    differential equation as discussed below. Has the same structure
                    and shape as `state`.
                  variables: represent the solution of the adjoint equation for
                    variable gradients. Represented as a `Tuple(Tensor, ...)` with as
                    many tensors as there are `variables` variable outside this
                    function.
                  constants: represent the solution of the adjoint equation for
                    constant gradients. Has the same structure and shape as
                    `constants` variable outside this function.
  
                  The adjoint sensitivity equation describes the gradient of the
                  solution with respect to the value of the solution at a previous
                  time t. Its dynamics are given by
                  d/dt[adj(t)] = -1 * adj(t) @ jacobian(ode_fn(t, z), z)
                  Which is computed as:
                  d/dt[adj(t)]_i = -1 * sum_j(adj(t)_j * d/dz_i[ode_fn(t, z)_j)]
                  d/dt[adj(t)]_i = -1 * d/dz_i[sum_j(no_grad_adj_j * ode_fn(t, z)_j)]
                  where in the last line we moved adj(t)_j under derivative by
                  removing gradient from it.
  
                  Adjoint equation for the gradient with respect to every
                  `tf.Variable` and constant theta follows:
                  d/dt[grad_theta(t)] = -1 * adj(t) @ jacobian(ode_fn(t, z), theta)
                  = -1 * d/d theta_i[sum_j(no_grad_adj_j * ode_fn(t, z)_j)]
  
                  Args:
                    backward_time: Floating `Tensor` representing current time.
                    augmented_state: `Tuple(state, adjoint_state, variable_grads)`
  
                  Returns:
                    negative_derivatives: Structure of `Tensor`s equal to backwards
                      time derivative of the `state` componnent.
                    adjoint_ode: Structure of `Tensor`s equal to backwards time
                      derivative of the `adjoint_state` component.
                    adjoint_variables_ode: Structure of `Tensor`s equal to backwards
                      time derivative of the `vars` component.
                    adjoint_constants_ode: Structure of `Tensor`s equal to backwards
                      time derivative of the `constants` component.
                  """
                  # The negative signs disappears after the change of variables.
                  # The ODE solver cannot handle the case initial_time > final_time
                  # and hence a change of variables backward_time = -time is used.
                  time = -backward_time
                  state, adjoint_state, _, _ = augmented_state
  
                  # TODO(b/152464477): Doesn't work reliably in TF1.
                  def grad_fn(state, variables, constants):
                      del variables  # We compute these gradients via the GradientTape
                      # capturing them.
                      derivatives = ode_fn(time, state, **constants)
                      adjoint_no_grad = tf.nest.map_structure(
                          tf.stop_gradient, adjoint_state
                      )
                      negative_derivatives = rk_util.weighted_sum(
                          [-1.0], [derivatives]
                      )
  
                      def dot_prod(tensor_a, tensor_b):
                          return tf.reduce_sum(tensor_a * tensor_b)
  
                      # See docstring for details.
                      adjoint_dot_derivatives = tf.nest.map_structure(
                          dot_prod, adjoint_no_grad, derivatives
                      )
                      adjoint_dot_derivatives = tf.squeeze(
                          tf.add_n(
                              tf.nest.flatten(adjoint_dot_derivatives)
                          )
                      )
                      return (
                          adjoint_dot_derivatives,
                          negative_derivatives,
                      )
  
                  values = (state, tuple(variables), constants)
                  (
                      (_, negative_derivatives),
                      gradients,
                  ) = tfp_gradient.value_and_gradient(
                      grad_fn,
                      values,
                      has_aux=True,
                      use_gradient_tape=True,
                  )
  
                  (
                      adjoint_ode,
                      adjoint_variables_ode,
                      adjoint_constants_ode,
                  ) = tf.nest.map_structure(
                      lambda v, g: tf.zeros_like(v) if g is None else g,
                      values,
                      tuple(gradients),
                  )
                  return (
                      negative_derivatives,
                      adjoint_ode,
                      adjoint_variables_ode,
                      adjoint_constants_ode,
                  )
  
              def make_augmented_state(n, prev_augmented_state):
                  """Constructs the augmented state for step `n`."""
                  (
                      _,
                      adjoint_state,
                      adjoint_variable_state,
                      adjoint_constant_state,
                  ) = prev_augmented_state
                  initial_state = _read_solution_components(
                      result_state_arrays,
                      input_state_structure,
                      n - 1,
                  )
                  initial_adjoint = _read_solution_components(
                      dresult_state_arrays,
                      input_state_structure,
                      n - 1,
                  )
                  initial_adjoint_state = rk_util.weighted_sum(
                      [1.0, 1.0], [adjoint_state, initial_adjoint]
                  )
                  augmented_state = (
                      initial_state,
                      initial_adjoint_state,
                      adjoint_variable_state,
                      adjoint_constant_state,
                  )
                  return augmented_state
  
              def reverse_to_result_time(
                  n, augmented_state, solver_internal_state, _
              ):
                  """Integrates the augmented system backwards in time."""
                  lower_bound_of_integration = result_time_array.read(n)
                  upper_bound_of_integration = result_time_array.read(
                      n - 1
                  )
                  initial_augmented_state = make_augmented_state(
                      n, augmented_state
                  )
                  # TODO(b/138304303): Allow the user to specify the Hessian of
                  # `ode_fn` so that we can get the Jacobian of the adjoint system.
                  # TODO(b/143624114): Support higher order derivatives.
                  solver_internal_state = adjoint_solver._adjust_solver_internal_state_for_state_jump(  # pylint: disable=protected-access
                      ode_fn=augmented_ode_fn,
                      initial_time=-lower_bound_of_integration,
                      initial_state=initial_augmented_state,
                      previous_solver_internal_state=solver_internal_state,
                      previous_state=augmented_state,
                  )
                  augmented_results = adjoint_solver.solve(
                      ode_fn=augmented_ode_fn,
                      initial_time=-lower_bound_of_integration,
                      initial_state=initial_augmented_state,
                      solution_times=[-upper_bound_of_integration],
                      batch_ndims=batch_ndims,
                      previous_solver_internal_state=solver_internal_state,
                  )
                  # Results added an extra time dim of size 1, squeeze it.
                  select_result = lambda x: tf.squeeze(x, [0])
                  result_state = augmented_results.states
                  result_state = tf.nest.map_structure(
                      select_result, result_state
                  )
                  status = augmented_results.diagnostics.status
                  return (
                      n - 1,
                      result_state,
                      augmented_results.solver_internal_state,
                      status,
                  )
  
              initial_n = num_result_times - 1
              solver_internal_state = adjoint_solver._initialize_solver_internal_state(  # pylint: disable=protected-access
                  ode_fn=augmented_ode_fn,
                  initial_time=result_time_array.read(initial_n),
                  initial_state=make_augmented_state(
                      initial_n, terminal_augmented_state
                  ),
              )
  
              _, augmented_state, _, _ = tf.while_loop(
                  lambda n, _as, _sis, status: (n >= 1)
                  & tf.equal(status, 0),
                  reverse_to_result_time,
                  (
                      initial_n,
                      terminal_augmented_state,
                      solver_internal_state,
                      0,
                  ),
                  back_prop=False,
              )
              (
                  _,
                  adjoint_state,
                  adjoint_variables,
                  adjoint_constants,
              ) = augmented_state
  
              if variables:
                  return (adjoint_state, adjoint_constants), list(
                      adjoint_variables
                  )
              else:
                  return adjoint_state, adjoint_constants
  
      @tfp_custom_gradient.custom_gradient(
          vjp_fwd=vjp_fwd,
          vjp_bwd=vjp_bwd,
      )
      def gradient_helper(initial_state, constants):
          """Restricts gradient to initial state components and constants."""
          return vjp_fwd(initial_state, constants)[0]
  
      # TODO(b/140760650): We must use a resource-using variable scope, otherwise
      # custom_gradient will complain even if there are no variables in `ode_fn`.
      with tf1.variable_scope(
          tf1.get_variable_scope(), use_resource=True
      ):
          return gradient_helper(initial_state, constants)
  
imports: import abc
  import collections
  import functools
  import six
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import custom_gradient
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.math import gradient
  from tensorflow_probability.python.math.ode import runge_kutta_util
  from tensorflow_probability.python.math.ode import util
label_desc: Use image node_id 3 for calling the Solver obj's underlying member method code with example usage: obj.solve(ode_fn, initial_time, initial_state, solution_times, jacobian_fn, jacobian_sparsity, batch_ndims, previous_solver_internal_state, constants) and returns: results,  results,  constants, vjp_fwd, gradient_helper, negative_derivatives,  adjoint_ode,  adjoint_variables_ode,  adjoint_constants_ode, augmented_state, unknown,  result_state,  augmented_results,  status, adjoint_state,  adjoint_constants,  list, adjoint_state,  adjoint_constants, adjoint_dot_derivatives,  negative_derivatives, tf