code: def interpolate(
      x: types.RealTensor,
      spline_data: SplineParameters,
      optimize_for_tpu: bool = False,
      dtype: tf.DType = None,
      name: str = None,
  ) -> types.RealTensor:
      """Interpolates spline values for the given `x` and the `spline_data`.
  
      Constant extrapolation is performed for the values outside the domain
      `spline_data.x_data`. This means that for `x > max(spline_data.x_data)`,
      `interpolate(x, spline_data) = spline_data.y_data[-1]`
      and for  `x < min(spline_data.x_data)`,
      `interpolate(x, spline_data) = spline_data.y_data[0]`.
  
      For the interpolation formula refer to p.548 of [1].
  
      #### References:
      [1]: R. Sedgewick, Algorithms in C, 1990, p. 545-550.
        Link: http://index-of.co.uk/Algorithms/Algorithms%20in%20C.pdf
  
      Args:
        x: A real `Tensor` of shape `batch_shape + [num_points]`.
        spline_data: An instance of `SplineParameters`. `spline_data.x_data` should
          have the same batch shape as `x`.
        optimize_for_tpu: A Python bool. If `True`, the algorithm uses one-hot
          encoding to lookup indices of `x` in `spline_data.x_data`. This
          significantly improves performance of the algorithm on a TPU device but
          may slow down performance on the CPU.
          Default value: `False`.
        dtype: Optional dtype for `x`.
          Default value: `None` which maps to the default dtype inferred by
            TensorFlow.
        name: Python `str` name prefixed to ops created by this function.
          Default value: `None` which is mapped to the default name
            `cubic_spline_interpolate`.
  
      Returns:
          A `Tensor` of the same shape and `dtype` as `x`. Represents
          the interpolated values.
  
      Raises:
        ValueError:
          If `x` batch shape is different from `spline_data.x_data` batch
          shape.
      """
      name = name or "cubic_spline_interpolate"
      with tf.name_scope(name):
          x = tf.convert_to_tensor(x, dtype=dtype, name="x")
          dtype = x.dtype
          # Unpack the spline data
          x_data = spline_data.x_data
          y_data = spline_data.y_data
          spline_coeffs = spline_data.spline_coeffs
          # Try broadcast batch_shapes
          (
              x,
              x_data,
              y_data,
              spline_coeffs,
          ) = tff_utils.broadcast_common_batch_shape(
              x, x_data, y_data, spline_coeffs
          )
          # Determine the splines to use.
          indices = tf.searchsorted(x_data, x, side="right") - 1
          # This selects all elements for the start of the spline interval.
          # Make sure indices lie in the permissible range
          lower_encoding = tf.maximum(indices, 0)
          # This selects all elements for the end of the spline interval.
          # Make sure indices lie in the permissible range
          upper_encoding = tf.minimum(
              indices + 1, tff_utils.get_shape(x_data)[-1] - 1
          )
          # Prepare indices for `tf.gather` or `tf.one_hot`
          # TODO(b/156720909): Extract get_slice logic into a common utilities module
          # for cubic and linear interpolation
          if optimize_for_tpu:
              x_data_size = tff_utils.get_shape(x_data)[-1]
              lower_encoding = tf.one_hot(
                  lower_encoding, x_data_size, dtype=dtype
              )
              upper_encoding = tf.one_hot(
                  upper_encoding, x_data_size, dtype=dtype
              )
  
          # Calculate dx and dy.
          # Simplified logic:
          # dx = x_data[indices + 1] - x_data[indices]
          # dy = y_data[indices + 1] - y_data[indices]
          # indices is a tensor with different values per row/spline
          def get_slice(x, encoding):
              if optimize_for_tpu:
                  return tf.math.reduce_sum(
                      tf.expand_dims(x, axis=-2) * encoding, axis=-1
                  )
              else:
                  return tf.gather(
                      x, encoding, axis=-1, batch_dims=x.shape.rank - 1
                  )
  
          x0 = get_slice(x_data, lower_encoding)
          x1 = get_slice(x_data, upper_encoding)
          dx = x1 - x0
  
          y0 = get_slice(y_data, lower_encoding)
          y1 = get_slice(y_data, upper_encoding)
          dy = y1 - y0
  
          spline_coeffs0 = get_slice(spline_coeffs, lower_encoding)
          spline_coeffs1 = get_slice(spline_coeffs, upper_encoding)
  
          t = (x - x0) / dx
          t = tf.where(dx > 0, t, tf.zeros_like(t))
          df = ((t + 1.0) * spline_coeffs1 * 2.0) - (
              (t - 2.0) * spline_coeffs0 * 2.0
          )
          df1 = df * t * (t - 1) / 6.0
          result = y0 + (t * dy) + (dx * dx * df1)
          # Use constant extrapolation outside the domain
          upper_bound = tf.expand_dims(
              tf.reduce_max(x_data, -1), -1
          ) + tf.zeros_like(result)
          lower_bound = tf.expand_dims(
              tf.reduce_min(x_data, -1), -1
          ) + tf.zeros_like(result)
          result = tf.where(
              tf.logical_and(x <= upper_bound, x >= lower_bound),
              result,
              tf.where(x > upper_bound, y0, y1),
          )
          return result
  
file_path: tf-quant-finance/tf_quant_finance/math/interpolation/cubic/cubic_interpolation.py
imports: import enum
  import tensorflow.compat.v2
  from tf_quant_finance import types
  from tf_quant_finance import utils
label_desc: Use image node_id 2 for calling a global function with example usage: interpolate(x, spline_data, optimize_for_tpu, dtype, name) and returns: result, tf, tf