class_docstr: Dormand-Prince explicit solver for non-stiff ODEs.
  
  Implements 5th order Runge-Kutta with adaptive step size control
  and dense output, using the Dormand-Prince method. Similar to the 'dopri5'
  method of `scipy.integrate.ode` and MATLAB's `ode45`. For details see [1].
  For solver API see `tfp.math.ode.Solver`.
  
  #### References
  
  [1]: Shampine, L. F. (1986). Some practical runge-kutta formulas.
      Mathematics of Computation, 46(173), 135-150, doi:10.2307/2008219
class_docstr_tok: ["Dormand-Prince","explicit","solver","for","non-stiff","ODEs",".","Implements","5th","order","Runge-Kutta","with","adaptive","step","size","control","and","dense","output",",","using","the","Dormand-Prince","method",".","Similar","to","the","'dopri5'","method","of","`","scipy.integrate.ode","`","and","MATLAB","'s","`","ode45","`",".","For","details","see","[","1","]",".","For","solver","API","see","`","tfp.math.ode.Solver","`",".","#","#","#","#","References","[","1","]",":","Shampine",",","L.","F.","(","1986",")",".","Some","practical","runge-kutta","formulas",".","Mathematics","of","Computation",",","46","(","173",")",",","135-150",",","doi:10.2307\/2008219"]
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,
  ):
      # Static assertions
      del jacobian_fn, jacobian_sparsity  # not used by DormandPrince
      if batch_ndims is not None and batch_ndims != 0:
          raise NotImplementedError(
              "For homogeneous batching use `batch_ndims=0`."
          )
      solution_times_by_solver = isinstance(
          solution_times, base.ChosenBySolver
      )
  
      with tf.name_scope(self._name):
          # (2) Convert to tensors, determine dtypes.
          p = self._prepare_common_params(initial_state, initial_time)
  
          max_num_steps = self._max_num_steps
          max_ode_fn_evals = 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_ode_fn_evals = (
                  max_num_steps * self.ODE_FN_EVALS_PER_STEP
              )
          step_size = tf.convert_to_tensor(
              self._first_step_size, dtype=p.real_dtype
          )
          rtol = tf.convert_to_tensor(tf.cast(self._rtol, p.real_dtype))
          atol = tf.convert_to_tensor(tf.cast(self._atol, p.real_dtype))
          safety = tf.convert_to_tensor(
              self._safety_factor, dtype=p.real_dtype
          )
          # Use i(d)factor notation for increasing and decreasing factors.
          ifactor, dfactor = (
              self._max_step_size_factor,
              self._min_step_size_factor,
          )
          ifactor = tf.convert_to_tensor(ifactor, dtype=p.real_dtype)
          dfactor = tf.convert_to_tensor(dfactor, dtype=p.real_dtype)
  
          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=p.initial_state,
                      initial_time=p.initial_time,
                  )
              )
  
          num_solution_times = 0
          if solution_times_by_solver:
              final_time = tf.cast(
                  solution_times.final_time, p.real_dtype
              )
              times_array = tf.TensorArray(
                  p.real_dtype,
                  size=num_solution_times,
                  dynamic_size=True,
                  element_shape=tf.TensorShape([]),
              )
          else:
              solution_times = tf.cast(solution_times, p.real_dtype)
              util.error_if_not_vector(solution_times, "solution_times")
              num_solution_times = ps.size(solution_times)
              times_array = tf.TensorArray(
                  p.real_dtype,
                  size=num_solution_times,
                  dynamic_size=False,
                  element_shape=[],
              ).unstack(solution_times)
  
          solutions_arrays = nest.map_structure_up_to(
              p.state_dtypes,
              lambda shape, dtype: tf.TensorArray(  # pylint: disable=g-long-lambda
                  dtype=dtype,
                  size=num_solution_times,
                  element_shape=shape,
                  dynamic_size=solution_times_by_solver,
              ),
              p.state_shapes,
              p.state_dtypes,
          )
  
          rk_step = functools.partial(
              self._step,
              max_ode_fn_evals=max_ode_fn_evals,
              ode_fn=ode_fn,
              atol=atol,
              rtol=rtol,
              safety=safety,
              ifactor=ifactor,
              dfactor=dfactor,
          )
          advance_to_solution_time = functools.partial(
              _advance_to_solution_time,
              times_array=solution_times,
              step_fn=rk_step,
              validate_args=self._validate_args,
          )
  
          assert_ops = self._assert_ops(
              ode_fn=ode_fn,
              initial_time=p.initial_time,
              initial_state=p.initial_state,
              solution_times=solution_times,
              previous_solver_state=previous_solver_internal_state,
              rtol=rtol,
              atol=atol,
              first_step_size=step_size,
              safety_factor=safety,
              min_step_size_factor=ifactor,
              max_step_size_factor=dfactor,
              max_num_steps=max_num_steps,
              solution_times_by_solver=solution_times_by_solver,
          )
          with tf.control_dependencies(assert_ops):
              ode_evals_by_now = 1 if self._validate_args else 0
              ode_evals_by_now += (
                  1 if solver_internal_state is None else 0
              )
              diagnostics = _DopriDiagnostics(
                  num_ode_fn_evaluations=ode_evals_by_now,
                  num_jacobian_evaluations=0,
                  num_matrix_factorizations=0,
                  status=0,
              )
  
              if solution_times_by_solver:
                  r = _dense_solutions_to_final_time(
                      final_time=final_time,
                      solver_state=solver_internal_state,
                      diagnostics=diagnostics,
                      step_fn=rk_step,
                      ode_fn=ode_fn,
                      times_array=times_array,
                      solutions_arrays=solutions_arrays,
                      validate_args=self._validate_args,
                  )
                  (
                      solver_internal_state,
                      diagnostics,
                      times_array,
                      solutions_arrays,
                  ) = r
              else:
  
                  def iterate_cond(time_id, *_):
                      return time_id < num_solution_times
  
                  [
                      _,
                      solver_internal_state,
                      diagnostics,
                      solutions_arrays,
                  ] = tf.while_loop(
                      iterate_cond,
                      advance_to_solution_time,
                      [
                          0,
                          solver_internal_state,
                          diagnostics,
                          solutions_arrays,
                      ],
                  )
  
              times = times_array.stack()
              stack_components = lambda x: x.stack()
              states = tf.nest.map_structure(
                  stack_components, solutions_arrays
              )
              return base.Results(
                  times=times,
                  states=states,
                  diagnostics=diagnostics,
                  solver_internal_state=solver_internal_state,
              )
  
imports: import collections
  import functools
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.math.ode import base
  from tensorflow_probability.python.math.ode import runge_kutta_util
  from tensorflow_probability.python.math.ode import util
  from tensorflow.python.util import nest
label_desc: Use image node_id 2 for calling the DormandPrince 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: base, unknown