code: def _prepare_grid(times, times_grid, *params):
      """Prepares grid of times for path generation.
  
      Args:
        times:  Rank 1 `Tensor` of increasing positive real values. The times at
          which the path points are to be evaluated.
        times_grid: An optional rank 1 `Tensor` representing time discretization
          grid. If `times` are not on the grid, then the nearest points from the
          grid are used.
        *params: Parameters of the Heston model. Either scalar `Tensor`s of the
          same `dtype` or instances of `PiecewiseConstantFunc`.
  
      Returns:
        Tuple `(all_times, mask)`.
        `all_times` is a 1-D real `Tensor` containing all points from 'times`, the
        uniform grid of points between `[0, times[-1]]` with grid size equal to
        `time_step`, and jump locations of piecewise constant parameters The
        `Tensor` is sorted in ascending order and may contain duplicates.
        `mask` is a boolean 1-D `Tensor` of the same shape as 'all_times', showing
        which elements of 'all_times' correspond to THE values from `times`.
        Guarantees that times[0]=0 and mask[0]=False.
      """
      if times_grid is None:
          additional_times = []
          for param in params:
              if hasattr(param, "is_piecewise_constant"):
                  if param.is_piecewise_constant:
                      # Flatten all jump locations
                      additional_times.append(
                          tf.reshape(param.jump_locations(), [-1])
                      )
          zeros = tf.constant([0], dtype=times.dtype)
          all_times = tf.concat(
              [zeros] + [times] + additional_times, axis=0
          )
          all_times = tf.sort(all_times)
          time_indices = tf.searchsorted(
              all_times, times, out_type=tf.int32
          )
      else:
          all_times = times_grid
          time_indices = tf.searchsorted(
              times_grid, times, out_type=tf.int32
          )
          # Adjust indices to bring `times` closer to `times_grid`.
          times_diff_1 = tf.gather(times_grid, time_indices) - times
          times_diff_2 = (
              tf.gather(times_grid, tf.nn.relu(time_indices - 1))
              - times
          )
          time_indices = tf.where(
              tf.math.abs(times_diff_2) > tf.math.abs(times_diff_1),
              time_indices,
              tf.nn.relu(time_indices - 1),
          )
      # Create a boolean mask to identify the iterations that have to be recorded.
      mask = tf.scatter_nd(
          indices=tf.expand_dims(
              tf.cast(time_indices, dtype=tf.int64), axis=1
          ),
          updates=tf.fill(tf.shape(times), True),
          shape=tf.shape(all_times, out_type=tf.int64),
      )
      return all_times, mask
  
imports: from typing import Callable, Union, Tuple, Optional
  import tensorflow.compat.v2
  from tf_quant_finance import types
  from tf_quant_finance import utils
  from tf_quant_finance.math import gradient
  from tf_quant_finance.math import piecewise
  from tf_quant_finance.math import random
  from tf_quant_finance.models import euler_sampling
  from tf_quant_finance.models import generic_ito_process
  from tf_quant_finance.models import utils
label_desc: Use image node_id 2 for calling a global function with example usage: _prepare_grid(times, times_grid) and returns: all_times,  mask