code: def _piecewise_constant_integrate(
      x1, x2, jump_locations, values, batch_rank
  ):
      """Integrates piecewise constant function between `x1` and `x2`."""
      # Initializer already verified that `jump_locations` and `values` have the
      # same shape.
      # Expand batch size to one if there is no batch shape.
      if x1.shape.as_list()[:batch_rank]:
          no_batch_shape = False
      else:
          no_batch_shape = True
          x1 = tf.expand_dims(x1, 0)
          x2 = tf.expand_dims(x2, 0)
      if not jump_locations.shape.as_list()[:-1]:
          jump_locations = tf.expand_dims(jump_locations, 0)
          values = tf.expand_dims(values, 0)
          batch_rank += 1
      # Compute integral values between the jump locations
      event_shape = tf.shape(values)[(batch_rank + 1) :]
      event_rank = values.shape.rank - batch_rank - 1
      num_data_points = tf.shape(values)[batch_rank]
      diff = jump_locations[..., 1:] - jump_locations[..., :-1]
      # Broadcast `diff` to the shape of
      # `batch_shape + [num_data_points - 2] + [1] * sample_rank`.
      for _ in range(event_rank):
          diff = tf.expand_dims(diff, -1)
      slice_indices = batch_rank * [slice(None)]
      slice_indices += [slice(1, num_data_points - 1)]
      integrals = tf.cumsum(values[slice_indices] * diff, batch_rank)
      # Pad integrals with zero values on left and right.
      batch_shape = tf.shape(integrals)[:batch_rank]
      pad_shape = tf.concat([batch_shape, [1], event_shape], axis=0)
      zeros = tf.zeros(pad_shape, dtype=integrals.dtype)
      integrals = tf.concat([zeros, integrals, zeros], axis=batch_rank)
      # Get jump locations and values and the integration end points
      value1, jump_location1, indices_1 = _get_indices_and_values(
          x1, jump_locations, values, "left", batch_rank
      )
      value2, jump_location2, indices_2 = _get_indices_and_values(
          x2, jump_locations, values, "right", batch_rank
      )
      integrals1 = tf.gather(
          integrals, indices_1, axis=batch_rank, batch_dims=batch_rank
      )
      integrals2 = tf.gather(
          integrals, indices_2, axis=batch_rank, batch_dims=batch_rank
      )
      # Broadcast `x1`, `x2`, `jump_location1`, `jump_location2` to the shape
      # `batch_shape + [num_points] + [1] * sample_rank`.
      for _ in range(event_rank):
          x1 = tf.expand_dims(x1, -1)
          x2 = tf.expand_dims(x2, -1)
          jump_location1 = tf.expand_dims(jump_location1, -1)
          jump_location2 = tf.expand_dims(jump_location2, -1)
      # Compute the value of the integrals.
      res = (
          (jump_location1 - x1) * value1
          + (x2 - jump_location2) * value2
          + integrals2
          - integrals1
      )
      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 3 for calling a global function with example usage: _piecewise_constant_integrate(x1, x2, jump_locations, values, batch_rank) and returns: tf, res