code: def _get_time_steps_info(
      start_time, end_time, num_steps, time_step, time_direction_fn
  ):
      """Creates a callable to step through time and estimates the max steps."""
      # time_direction_fn must be one of _time_step_forward_fn and
      # _time_step_backward_fn
      dt = None
      estimated_max_steps = None
      interval = tf.math.abs(end_time - start_time)
      if num_steps is not None:
          dt = interval / tf.cast(num_steps, dtype=start_time.dtype)
          estimated_max_steps = num_steps
      if time_step is not None and not _is_callable(time_step):
          dt = time_step
          estimated_max_steps = tf.cast(
              tf.math.ceil(interval / dt), dtype=tf.int32
          )
      if dt is not None:
          raw_time_step_fn = lambda _: dt
      else:
          raw_time_step_fn = time_step
  
      def time_step_fn(t):
          # t is the current time.
          # t_next is the next time
          dt = raw_time_step_fn(t)
          should_stop, t_next = time_direction_fn(t, dt, end_time)
          return should_stop, t_next
  
      return time_step_fn, estimated_max_steps
  
file_path: tf-quant-finance/tf_quant_finance/math/pde/fd_solvers.py
imports: import tensorflow.compat.v2
  from tf_quant_finance import utils
  from tf_quant_finance.math.pde.steppers.douglas_adi import douglas_adi_step
  from tf_quant_finance.math.pde.steppers.oscillation_damped_crank_nicolson import oscillation_damped_crank_nicolson_step
label_desc: Use image node_id 5 for calling a global function with example usage: _get_time_steps_info(start_time, end_time, num_steps, time_step, time_direction_fn) and returns: time_step_fn,  estimated_max_steps, should_stop,  t_next