code: def runge_kutta_step(
      ode_fn, y0, f0, t0, dt, tableau, name="runge_kutta_step"
  ):
      """Take an arbitrary Runge-Kutta step and estimate error.
  
      Args:
        ode_fn: Callable(t, y) -> dy_dt that evaluate the time derivative of `y`.
        y0: `Tensor` initial value for the state.
        f0: `Tensor` initial value for the derivative of `y0` = `ode_fn(t0, y0)`.
        t0: `Tensor` value for the initial time.
        dt: `Tensor` value for the desired time step.
        tableau: `ButcherTableau` describing how to take the Runge-Kutta step.
        name: optional name for the operation.
  
      Returns:
        rk_state_tuple: Tuple `(y1, f1, y1_error, k)` giving the estimated function
          value after the Runge-Kutta step at `t1 = t0 + dt`, the derivative of the
          state at `t1`, estimated error at `t1`, and a list of Runge-Kutta
          coefficients `k` used for calculating these terms.
      """
      with tf.name_scope(name):
          y0 = tf.nest.map_structure(tf.convert_to_tensor, y0)
          f0 = tf.nest.map_structure(tf.convert_to_tensor, f0)
          t0 = tf.convert_to_tensor(t0, name="t0")
          dt = tf.convert_to_tensor(dt, name="dt")
  
          k = [f0]
          for a_i, b_i in zip(tableau.a, tableau.b):
              ti = t0 + a_i * dt
              yi = weighted_sum([1.0, dt], [y0, weighted_sum(b_i, k)])
              k.append(ode_fn(ti, yi))
  
          if not (
              tableau.c_sol[-1] == 0
              and tableau.c_sol[:-1] == tableau.b[-1]
          ):
              # This property (true for Dormand-Prince) lets us save a few FLOPs.
              yi = weighted_sum(
                  [1.0, dt], [y0, weighted_sum(tableau.c_sol, k)]
              )
  
          y1 = yi
          f1 = k[-1]
          y1_error = weighted_sum(
              [dt], [weighted_sum(tableau.c_error, k)]
          )
          return y1, f1, y1_error, k
  
file_path: probability/tensorflow_probability/python/math/ode/runge_kutta_util.py
imports: import collections
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
label_desc: Use image node_id 10 for calling a global function with example usage: runge_kutta_step(ode_fn, y0, f0, t0, dt, tableau, name) and returns: y1,  f1,  y1_error,  k