code: def _piecewise_constant_function(
      x, jump_locations, values, batch_rank, side="left"
  ):
      """Computes value of the piecewise constant function."""
      # Initializer already verified that `jump_locations` and `values` have the
      # same shape
      batch_shape = jump_locations.shape.as_list()[:-1]
      # Check that the batch shape of `x` is the same as of `jump_locations` and
      # `values`
      batch_shape_x = x.shape.as_list()[:batch_rank]
      if batch_shape_x != batch_shape:
          raise ValueError(
              "Batch shape of `x` is {1} but should be {0}".format(
                  batch_shape, batch_shape_x
              )
          )
      if x.shape.as_list()[:batch_rank]:
          no_batch_shape = False
      else:
          no_batch_shape = True
          x = tf.expand_dims(x, 0)
          batch_rank += 1
      # Expand batch size to one if there is no batch shape
      if not batch_shape:
          jump_locations = tf.expand_dims(jump_locations, 0)
          values = tf.expand_dims(values, 0)
      indices = tf.searchsorted(jump_locations, x, side=side)
      res = tf.gather(
          values, indices, axis=batch_rank, batch_dims=batch_rank
      )
      if no_batch_shape:
          return tf.squeeze(res, 0)
      else:
          return res
  
file_path: tf-quant-finance/tf_quant_finance/math/piecewise.py
imports: import tensorflow.compat.v2
label_desc: Use image node_id 2 for calling a global function with example usage: _piecewise_constant_function(x, jump_locations, values, batch_rank, side) and returns: tf, res