code: def nest_where(accept_step, new_values, old_values):
      """Returns `new_values` if `accept_step` is True `old_values` otherwise.
  
      Uses `tf.where` on individual elements to select `new_values` or `old_values`.
  
      Args:
        accept_step: Scalar boolean `Tensor` indicating whether to return
          `new_values`.
        new_values: Possible nested structure of `Tensor`s.
        old_values: Possible nested structure of `Tensor`s. Must have the same
          structure as `new_values`.
  
      Returns:
        values: `new_values` if `accept_step` is True and `old_values` otherwise.
      """
      tf.nest.assert_same_structure(new_values, old_values)
      select_new_or_old = lambda x, y: tf.where(accept_step, x, y)
      values = tf.nest.map_structure(
          select_new_or_old, new_values, old_values
      )
      return values
  
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 6 for calling a global function with example usage: nest_where(accept_step, new_values, old_values) and returns: values