class_docstr: Backward differentiation formula (BDF) solver for stiff ODEs.
  
  Implements the solver described in [Shampine and Reichelt (1997)][1], a
  variable step size, variable order (VSVO) BDF integrator with order varying
  between 1 and 5.
  
  #### Algorithm details
  
  Each step involves solving the following nonlinear equation by Newton's
  method:
  ```none
  0 = 1/1 * BDF(1, y)[n+1] + ... + 1/k * BDF(k, y)[n+1]
    - h ode_fn(t[n+1], y[n+1])
    - bdf_coefficients[k-1] * (1/1 + ... + 1/k) * (y[n+1] - y[n] - BDF(1, y)[n]
                                                          -  ... - BDF(k, y)[n])
  ```
  where `k >= 1` is the current order of the integrator, `h` is the current step
  size, `bdf_coefficients` is a list of numbers that parameterizes the method,
  and `BDF(m, y)` is the `m`-th order backward difference of the vector `y`. In
  particular, `BDF(0, y)[n] = y[n]` and
  `BDF(m + 1, y)[n] = BDF(m, y)[n] - BDF(m, y)[n - 1]` for `m >= 0`.
  
  Newton's method can fail because
  * the method has exceeded the maximum number of iterations,
  * the method is converging too slowly, or
  * the method is not expected to converge
  (the last two conditions are determined by approximating the Lipschitz
  constant associated with the iteration).
  
  When `evaluate_jacobian_lazily` is `True`, the solver avoids evaluating the
  Jacobian of the dynamics function as much as possible. In particular, Newton's
  method will try to use the Jacobian from a previous integration step. If
  Newton's method fails with an out-of-date Jacobian, the Jacobian is
  re-evaluated and Newton's method is restarted. If Newton's method fails and
  the Jacobian is already up-to-date, then the step size is decreased and
  Newton's method is restarted.
  
  Even if Newton's method converges, the solution it generates can still be
  rejected if it exceeds the specified error tolerance due to truncation error.
  In this case, the step size is decreased and Newton's method is restarted.
  
  #### References
  
  [1]: Lawrence F. Shampine and Mark W. Reichelt. The MATLAB ODE Suite. _SIAM
       Journal on Scientific Computing_ 18(1):1-22, 1997.
class_docstr_tok: ["Backward","differentiation","formula","(","BDF",")","solver","for","stiff","ODEs",".","Implements","the","solver","described","in","[","Shampine","and","Reichelt","(","1997",")","]","[","1","]",",","a","variable","step","size",",","variable","order","(","VSVO",")","BDF","integrator","with","order","varying","between","1","and","5",".","#","#","#","#","Algorithm","details","Each","step","involves","solving","the","following","nonlinear","equation","by","Newton's","method",":","``","`","none","0","=","1\/1","*","BDF","(","1",",","y",")","[","n+1","]","+","...","+","1\/k","*","BDF","(","k",",","y",")","[","n+1","]","-","h","ode_fn","(","t","[","n+1","]",",","y","[","n+1","]",")","-","bdf_coefficients","[","k-1","]","*","(","1\/1","+","...","+","1\/k",")","*","(","y","[","n+1","]","-","y","[","n","]","-","BDF","(","1",",","y",")","[","n","]","-","...","-","BDF","(","k",",","y",")","[","n","]",")","``","`","where","`","k",">","=","1","`","is","the","current","order","of","the","integrator",",","`","h","`","is","the","current","step","size",",","`","bdf_coefficients","`","is","a","list","of","numbers","that","parameterizes","the","method",",","and","`","BDF","(","m",",","y",")","`","is","the","`","m","`","-th","order","backward","difference","of","the","vector","`","y","`",".","In","particular",",","`","BDF","(","0",",","y",")","[","n","]","=","y","[","n","]","`","and","`","BDF","(","m","+","1",",","y",")","[","n","]","=","BDF","(","m",",","y",")","[","n","]","-","BDF","(","m",",","y",")","[","n","-","1","]","`","for","`","m",">","=","0","`",".","Newton","'s","method","can","fail","because","*","the","method","has","exceeded","the","maximum","number","of","iterations",",","*","the","method","is","converging","too","slowly",",","or","*","the","method","is","not","expected","to","converge","(","the","last","two","conditions","are","determined","by","approximating","the","Lipschitz","constant","associated","with","the","iteration",")",".","When","`","evaluate_jacobian_lazily","`","is","`","True","`",",","the","solver","avoids","evaluating","the","Jacobian","of","the","dynamics","function","as","much","as","possible",".","In","particular",",","Newton's","method","will","try","to","use","the","Jacobian","from","a","previous","integration","step",".","If","Newton","'s","method","fails","with","an","out-of-date","Jacobian",",","the","Jacobian","is","re-evaluated","and","Newton","'s","method","is","restarted",".","If","Newton","'s","method","fails","and","the","Jacobian","is","already","up-to-date",",","then","the","step","size","is","decreased","and","Newton","'s","method","is","restarted",".","Even","if","Newton","'s","method","converges",",","the","solution","it","generates","can","still","be","rejected","if","it","exceeds","the","specified","error","tolerance","due","to","truncation","error",".","In","this","case",",","the","step","size","is","decreased","and","Newton","'s","method","is","restarted",".","#","#","#","#","References","[","1","]",":","Lawrence","F.","Shampine","and","Mark","W.","Reichelt",".","The","MATLAB","ODE","Suite",".","_SIAM","Journal","on","Scientific","Computing_","18","(","1",")",":1-22",",","1997","."]
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,
  ):
      # This function is comprised of the following sequential stages:
      # (1) Make static assertions.
      # (2) Initialize variables.
      # (3) Make non-static assertions.
      # (4) Solve up to final time.
      # (5) Return `Results` object.
      #
      # The stages can be found in the code by searching for (n) where n=1..5.
      #
      # By static vs. non-static assertions (see stages 1 and 3), we mean
      # assertions that can be made before the graph is run vs. those that can
      # only be made at run time. The latter are constructed as a list of
      # tf.Assert operations by the function `assert_ops` (see below).
      #
      # If `solution_times` is specified as a `Tensor`, stage 4 consists of three
      # nested loops, which can be conceptually understood as follows:
      # ```
      # current_time, current_state = initial_time, initial_state
      # order, step_size = 1, first_step_size
      # for solution_time in solution_times:
      #   while current_time < solution_time:
      #     while True:
      #       next_time = current_time + step_size
      #       next_state, error = (
      #           solve_nonlinear_equation_to_get_approximate_state_at_next_time(
      #           current_time, current_state, next_time, order))
      #       if error < tolerance:
      #         current_time, current_state = next_time, next_state
      #         order, step_size = (
      #           maybe_update_order_and_step_size(order, step_size))
      #         break
      #       else:
      #         step_size = decrease_step_size(step_size)
      # ```
      # The outermost loop advances the solver to the next `solution_time` (see
      # `advance_to_solution_time`). The middle loop advances the solver by a
      # small timestep (see `step`). The innermost loop determines the size of
      # that timestep (see `maybe_step`).
      #
      # If `solution_times` is specified as
      # `tfp.math.ode.ChosenBySolver(final_time)`, the outermost loop is skipped
      # and `solution_time` in the middle loop is replaced by `final_time`.
  
      def advance_to_solution_time(
          n,
          diagnostics,
          iterand,
          solver_internal_state,
          state_vec_array,
          time_array,
      ):
          """Takes multiple steps to advance time to `solution_times[n]`."""
  
          def step_cond(next_time, diagnostics, iterand, *_):
              return (iterand.time < next_time) & (
                  tf.equal(diagnostics.status, 0)
              )
  
          nth_solution_time = solution_time_array.read(n)
          [
              _,
              diagnostics,
              iterand,
              solver_internal_state,
              state_vec_array,
              time_array,
          ] = tf.while_loop(
              step_cond,
              step,
              [
                  nth_solution_time,
                  diagnostics,
                  iterand,
                  solver_internal_state,
                  state_vec_array,
                  time_array,
              ],
          )
          state_vec_array = state_vec_array.write(
              n, solver_internal_state.backward_differences[0]
          )
          time_array = time_array.write(n, nth_solution_time)
          return (
              n + 1,
              diagnostics,
              iterand,
              solver_internal_state,
              state_vec_array,
              time_array,
          )
  
      def step(
          next_time,
          diagnostics,
          iterand,
          solver_internal_state,
          state_vec_array,
          time_array,
      ):
          """Takes a single step."""
          distance_to_next_time = next_time - iterand.time
          overstepped = iterand.new_step_size > distance_to_next_time
          iterand = iterand._replace(
              new_step_size=tf1.where(
                  overstepped,
                  distance_to_next_time,
                  iterand.new_step_size,
              ),
              should_update_step_size=overstepped
              | iterand.should_update_step_size,
          )
  
          if not self._evaluate_jacobian_lazily:
              diagnostics = diagnostics._replace(
                  num_jacobian_evaluations=diagnostics.num_jacobian_evaluations
                  + 1
              )
              iterand = iterand._replace(
                  jacobian_mat=jacobian_fn_mat(
                      iterand.time,
                      solver_internal_state.backward_differences[0],
                  ),
                  jacobian_is_up_to_date=True,
              )
  
          def maybe_step_cond(accepted, diagnostics, *_):
              return tf.logical_not(accepted) & tf.equal(
                  diagnostics.status, 0
              )
  
          (
              _,
              diagnostics,
              iterand,
              solver_internal_state,
          ) = tf.while_loop(
              maybe_step_cond,
              maybe_step,
              [False, diagnostics, iterand, solver_internal_state],
          )
  
          if solution_times_chosen_by_solver:
              state_vec_array = state_vec_array.write(
                  state_vec_array.size(),
                  solver_internal_state.backward_differences[0],
              )
              time_array = time_array.write(
                  time_array.size(), iterand.time
              )
  
          return (
              next_time,
              diagnostics,
              iterand,
              solver_internal_state,
              state_vec_array,
              time_array,
          )
  
      def maybe_step(
          accepted, diagnostics, iterand, solver_internal_state
      ):
          """Takes a single step only if the outcome has a low enough error."""
          [
              num_jacobian_evaluations,
              num_matrix_factorizations,
              num_ode_fn_evaluations,
              status,
          ] = diagnostics
          [
              jacobian_mat,
              jacobian_is_up_to_date,
              new_step_size,
              num_steps,
              num_steps_same_size,
              should_update_jacobian,
              should_update_step_size,
              time,
              unitary,
              upper,
          ] = iterand
          [
              backward_differences,
              order,
              step_size,
          ] = solver_internal_state
  
          if max_num_steps is not None:
              status = tf1.where(
                  tf.equal(num_steps, max_num_steps), -1, 0
              )
  
          backward_differences = tf1.where(
              should_update_step_size,
              bdf_util.interpolate_backward_differences(
                  backward_differences, order, new_step_size / step_size
              ),
              backward_differences,
          )
          step_size = tf1.where(
              should_update_step_size, new_step_size, step_size
          )
          should_update_factorization = should_update_step_size
          num_steps_same_size = tf1.where(
              should_update_step_size, 0, num_steps_same_size
          )
  
          def update_factorization():
              return bdf_util.newton_qr(
                  jacobian_mat,
                  newton_coefficients_array.read(order),
                  step_size,
              )
  
          if self._evaluate_jacobian_lazily:
  
              def update_jacobian_and_factorization():
                  new_jacobian_mat = jacobian_fn_mat(
                      time, backward_differences[0]
                  )
                  new_unitary, new_upper = update_factorization()
                  return [
                      new_jacobian_mat,
                      True,
                      num_jacobian_evaluations + 1,
                      new_unitary,
                      new_upper,
                  ]
  
              def maybe_update_factorization():
                  new_unitary, new_upper = tf.cond(
                      should_update_factorization,
                      update_factorization,
                      lambda: [unitary, upper],
                  )
                  return [
                      jacobian_mat,
                      jacobian_is_up_to_date,
                      num_jacobian_evaluations,
                      new_unitary,
                      new_upper,
                  ]
  
              [
                  jacobian_mat,
                  jacobian_is_up_to_date,
                  num_jacobian_evaluations,
                  unitary,
                  upper,
              ] = tf.cond(
                  should_update_jacobian,
                  update_jacobian_and_factorization,
                  maybe_update_factorization,
              )
          else:
              unitary, upper = update_factorization()
              num_matrix_factorizations += 1
  
          tol = p.atol + p.rtol * tf.abs(backward_differences[0])
          newton_tol = newton_tol_factor * tf.norm(tol)
  
          [
              newton_converged,
              next_backward_difference,
              next_state_vec,
              newton_num_iters,
          ] = bdf_util.newton(
              backward_differences,
              max_num_newton_iters,
              newton_coefficients_array.read(order),
              p.ode_fn_vec,
              order,
              step_size,
              time,
              newton_tol,
              unitary,
              upper,
          )
          num_steps += 1
          num_ode_fn_evaluations += newton_num_iters
  
          # If Newton's method failed and the Jacobian was up to date, decrease the
          # step size.
          newton_failed = tf.logical_not(newton_converged)
          should_update_step_size = (
              newton_failed & jacobian_is_up_to_date
          )
          new_step_size = step_size * tf1.where(
              should_update_step_size, newton_step_size_factor, 1.0
          )
  
          # If Newton's method failed and the Jacobian was NOT up to date, update
          # the Jacobian.
          should_update_jacobian = newton_failed & tf.logical_not(
              jacobian_is_up_to_date
          )
  
          error_ratio = tf1.where(
              newton_converged,
              bdf_util.error_ratio(
                  next_backward_difference,
                  error_coefficients_array.read(order),
                  tol,
              ),
              np.nan,
          )
          accepted = error_ratio < 1.0
          converged_and_rejected = newton_converged & tf.logical_not(
              accepted
          )
  
          # If Newton's method converged but the solution was NOT accepted, decrease
          # the step size.
          new_step_size = tf1.where(
              converged_and_rejected,
              util.next_step_size(
                  step_size,
                  order,
                  error_ratio,
                  p.safety_factor,
                  min_step_size_factor,
                  max_step_size_factor,
              ),
              new_step_size,
          )
          should_update_step_size = (
              should_update_step_size | converged_and_rejected
          )
  
          # If Newton's method converged and the solution was accepted, update the
          # matrix of backward differences.
          time = tf1.where(accepted, time + step_size, time)
          backward_differences = tf1.where(
              accepted,
              bdf_util.update_backward_differences(
                  backward_differences,
                  next_backward_difference,
                  next_state_vec,
                  order,
              ),
              backward_differences,
          )
          jacobian_is_up_to_date = (
              jacobian_is_up_to_date & tf.logical_not(accepted)
          )
          num_steps_same_size = tf1.where(
              accepted, num_steps_same_size + 1, num_steps_same_size
          )
  
          # Order and step size are only updated if we have taken strictly more than
          # order + 1 steps of the same size. This is to prevent the order from
          # being throttled.
          should_update_order_and_step_size = accepted & (
              num_steps_same_size > order + 1
          )
  
          backward_differences_array = tf.TensorArray(
              backward_differences.dtype,
              size=bdf_util.MAX_ORDER + 3,
              clear_after_read=False,
              element_shape=next_backward_difference.shape,
          ).unstack(backward_differences)
          new_order = order
          new_error_ratio = error_ratio
          for offset in [-1, +1]:
              proposed_order = tf.clip_by_value(
                  order + offset, 1, max_order
              )
              proposed_error_ratio = bdf_util.error_ratio(
                  backward_differences_array.read(proposed_order + 1),
                  error_coefficients_array.read(proposed_order),
                  tol,
              )
              proposed_error_ratio_is_lower = (
                  proposed_error_ratio < new_error_ratio
              )
              new_order = tf1.where(
                  should_update_order_and_step_size
                  & proposed_error_ratio_is_lower,
                  proposed_order,
                  new_order,
              )
              new_error_ratio = tf1.where(
                  should_update_order_and_step_size
                  & proposed_error_ratio_is_lower,
                  proposed_error_ratio,
                  new_error_ratio,
              )
          order = new_order
          error_ratio = new_error_ratio
  
          new_step_size = tf1.where(
              should_update_order_and_step_size,
              util.next_step_size(
                  step_size,
                  order,
                  error_ratio,
                  p.safety_factor,
                  min_step_size_factor,
                  max_step_size_factor,
              ),
              new_step_size,
          )
          should_update_step_size = (
              should_update_step_size
              | should_update_order_and_step_size
          )
  
          diagnostics = _BDFDiagnostics(
              num_jacobian_evaluations,
              num_matrix_factorizations,
              num_ode_fn_evaluations,
              status,
          )
          iterand = _BDFIterand(
              jacobian_mat,
              jacobian_is_up_to_date,
              new_step_size,
              num_steps,
              num_steps_same_size,
              should_update_jacobian,
              should_update_step_size,
              time,
              unitary,
              upper,
          )
          solver_internal_state = _BDFSolverInternalState(
              backward_differences, order, step_size
          )
          return accepted, diagnostics, iterand, solver_internal_state
  
      # (1) Make static assertions.
      # TODO(b/138304296): Support specifying Jacobian sparsity patterns.
      if jacobian_sparsity is not None:
          raise NotImplementedError(
              "The BDF solver does not support specifying "
              "Jacobian sparsity patterns."
          )
      if batch_ndims is not None and batch_ndims != 0:
          raise NotImplementedError(
              "The BDF solver does not support batching."
          )
      solution_times_chosen_by_solver = isinstance(
          solution_times, base.ChosenBySolver
      )
  
      with tf.name_scope(self._name):
          # (2) Convert to tensors.
          p = self._prepare_common_params(
              ode_fn=ode_fn,
              initial_state=initial_state,
              initial_time=initial_time,
          )
  
          if jacobian_fn is None and dtype_util.is_complex(
              p.common_state_dtype
          ):
              raise NotImplementedError(
                  "The BDF solver does not support automatic "
                  "Jacobian computations for complex dtypes."
              )
  
          # Convert everything to operate on a single, concatenated vector form.
          jacobian_fn_mat = util.get_jacobian_fn_mat(
              jacobian_fn,
              p.ode_fn_vec,
              p.state_shape,
              dtype=p.common_state_dtype,
          )
  
          num_solution_times = 0
          if solution_times_chosen_by_solver:
              final_time = tf.cast(
                  solution_times.final_time, p.real_dtype
              )
          else:
              solution_times = tf.cast(solution_times, p.real_dtype)
              final_time = tf.reduce_max(solution_times)
              num_solution_times = ps.size(solution_times)
              solution_time_array = tf.TensorArray(
                  solution_times.dtype,
                  size=num_solution_times,
                  element_shape=[],
              ).unstack(solution_times)
              util.error_if_not_vector(solution_times, "solution_times")
          min_step_size_factor = tf.convert_to_tensor(
              self._min_step_size_factor, dtype=p.real_dtype
          )
          max_step_size_factor = tf.convert_to_tensor(
              self._max_step_size_factor, dtype=p.real_dtype
          )
          max_num_steps = self._max_num_steps
          if max_num_steps is not None:
              max_num_steps = tf.convert_to_tensor(
                  max_num_steps, dtype=tf.int32
              )
          max_order = tf.convert_to_tensor(
              self._max_order, dtype=tf.int32
          )
          max_num_newton_iters = self._max_num_newton_iters
          if max_num_newton_iters is not None:
              max_num_newton_iters = tf.convert_to_tensor(
                  max_num_newton_iters, dtype=tf.int32
              )
          newton_tol_factor = tf.convert_to_tensor(
              self._newton_tol_factor, dtype=p.real_dtype
          )
          newton_step_size_factor = tf.convert_to_tensor(
              self._newton_step_size_factor, dtype=p.real_dtype
          )
          (
              newton_coefficients,
              error_coefficients,
          ) = self._prepare_coefficients(p.common_state_dtype)
          if self._validate_args:
              final_time = tf.ensure_shape(final_time, [])
              min_step_size_factor = tf.ensure_shape(
                  min_step_size_factor, []
              )
              max_step_size_factor = tf.ensure_shape(
                  max_step_size_factor, []
              )
              if max_num_steps is not None:
                  max_num_steps = tf.ensure_shape(max_num_steps, [])
              max_order = tf.ensure_shape(max_order, [])
              if max_num_newton_iters is not None:
                  max_num_newton_iters = tf.ensure_shape(
                      max_num_newton_iters, []
                  )
              newton_tol_factor = tf.ensure_shape(newton_tol_factor, [])
              newton_step_size_factor = tf.ensure_shape(
                  newton_step_size_factor, []
              )
          newton_coefficients_array = tf.TensorArray(
              newton_coefficients.dtype,
              size=bdf_util.MAX_ORDER + 1,
              clear_after_read=False,
              element_shape=[],
          ).unstack(newton_coefficients)
          error_coefficients_array = tf.TensorArray(
              error_coefficients.dtype,
              size=bdf_util.MAX_ORDER + 1,
              clear_after_read=False,
              element_shape=[],
          ).unstack(error_coefficients)
          solver_internal_state = previous_solver_internal_state
          if solver_internal_state is None:
              solver_internal_state = (
                  self._initialize_solver_internal_state(
                      ode_fn=ode_fn,
                      initial_state=initial_state,
                      initial_time=initial_time,
                  )
              )
          state_vec_array = tf.TensorArray(
              p.common_state_dtype,
              size=num_solution_times,
              dynamic_size=solution_times_chosen_by_solver,
              element_shape=p.initial_state_vec.shape,
          )
          time_array = tf.TensorArray(
              p.real_dtype,
              size=num_solution_times,
              dynamic_size=solution_times_chosen_by_solver,
              element_shape=tf.TensorShape([]),
          )
          diagnostics = _BDFDiagnostics(
              num_jacobian_evaluations=0,
              num_matrix_factorizations=0,
              num_ode_fn_evaluations=0,
              status=0,
          )
          iterand = _BDFIterand(
              jacobian_mat=tf.zeros(
                  [p.num_odes, p.num_odes], dtype=p.common_state_dtype
              ),
              jacobian_is_up_to_date=False,
              new_step_size=solver_internal_state.step_size,
              num_steps=0,
              num_steps_same_size=0,
              should_update_jacobian=True,
              should_update_step_size=False,
              time=p.initial_time,
              unitary=tf.zeros(
                  [p.num_odes, p.num_odes], dtype=p.common_state_dtype
              ),
              upper=tf.zeros(
                  [p.num_odes, p.num_odes], dtype=p.common_state_dtype
              ),
          )
  
          # (3) Make non-static assertions.
          with tf.control_dependencies(
              self._assert_ops(
                  previous_solver_internal_state=previous_solver_internal_state,
                  initial_state_vec=p.initial_state_vec,
                  final_time=final_time,
                  initial_time=p.initial_time,
                  solution_times=solution_times,
                  max_num_steps=max_num_steps,
                  max_num_newton_iters=max_num_newton_iters,
                  atol=p.atol,
                  rtol=p.rtol,
                  first_step_size=solver_internal_state.step_size,
                  safety_factor=p.safety_factor,
                  min_step_size_factor=min_step_size_factor,
                  max_step_size_factor=max_step_size_factor,
                  max_order=max_order,
                  newton_tol_factor=newton_tol_factor,
                  newton_step_size_factor=newton_step_size_factor,
                  solution_times_chosen_by_solver=solution_times_chosen_by_solver,
              )
          ):
              # (4) Solve up to final time.
              if solution_times_chosen_by_solver:
  
                  def step_cond(next_time, diagnostics, iterand, *_):
                      return (iterand.time < next_time) & (
                          tf.equal(diagnostics.status, 0)
                      )
  
                  [
                      _,
                      diagnostics,
                      iterand,
                      solver_internal_state,
                      state_vec_array,
                      time_array,
                  ] = tf.while_loop(
                      step_cond,
                      step,
                      [
                          final_time,
                          diagnostics,
                          iterand,
                          solver_internal_state,
                          state_vec_array,
                          time_array,
                      ],
                  )
  
              else:
  
                  def advance_to_solution_time_cond(n, diagnostics, *_):
                      return (n < num_solution_times) & (
                          tf.equal(diagnostics.status, 0)
                      )
  
                  [
                      _,
                      diagnostics,
                      iterand,
                      solver_internal_state,
                      state_vec_array,
                      time_array,
                  ] = tf.while_loop(
                      advance_to_solution_time_cond,
                      advance_to_solution_time,
                      [
                          0,
                          diagnostics,
                          iterand,
                          solver_internal_state,
                          state_vec_array,
                          time_array,
                      ],
                  )
  
              # (6) Return `Results` object.
              states = util.get_state_from_vec(
                  state_vec_array.stack(), p.state_shape
              )
              times = time_array.stack()
              if not solution_times_chosen_by_solver:
                  tensorshape_util.set_shape(
                      times, solution_times.shape
                  )
                  tf.nest.map_structure(
                      lambda s, ini_s: tensorshape_util.set_shape(  # pylint: disable=g-long-lambda
                          s,
                          tensorshape_util.concatenate(
                              solution_times.shape, ini_s.shape
                          ),
                      ),
                      states,
                      p.initial_state,
                  )
              return base.Results(
                  times=times,
                  states=states,
                  diagnostics=diagnostics,
                  solver_internal_state=solver_internal_state,
              )
  
imports: import collections
  import functools
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.math.ode import base
  from tensorflow_probability.python.math.ode import bdf_util
  from tensorflow_probability.python.math.ode import util
label_desc: Use image node_id 2 for calling the BDF 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) and returns: unknown,  diagnostics,  iterand,  solver_internal_state,  state_vec_array,  time_array, next_time,  diagnostics,  iterand,  solver_internal_state,  state_vec_array,  time_array, accepted,  diagnostics,  iterand,  solver_internal_state, unknown, unknown, bdf_util, base, list, list, unknown, unknown