code: def segment_cumsum(
      x, segment_ids, exclusive=False, dtype=None, name=None
  ):
      """Computes cumulative sum of elements in a segment.
  
      For a complete description of segment_* ops see documentation of
      `tf.segment_sum`. This op extends the `tf.math.cumsum` functionality to
      segmented inputs.
  
      The behaviour of this op is the same as that of the op `tf.math.cumsum` within
      each segment. The result is effectively a concatenation of the results of
      `tf.math.cumsum` applied to each segment with the same interpretation for the
      argument `exclusive`.
  
      #### Example
  
      ```python
        x = tf.constant([2, 5, 1, 7, 9] + [32, 10, 12, 3] + [4, 8, 5])
        segments = tf.constant([0, 0, 0, 0, 0] + [1, 1, 1, 1] + [2, 2, 2])
        # Inclusive cumulative sum.
        # Expected result: [2, 7, 8, 15, 24, 32, 42, 54, 57, 4, 12, 17]
        cumsum1 = segment_cumsum(
            x, segment_ids=segments, exclusive=False)
        # Exclusive cumsum.
        # Expected result: [0, 2, 7, 8, 15, 0, 32, 42, 54, 0, 4, 12]
        cumsum2 = segment_cumsum(
            x, segment_ids=segments, exclusive=True)
      ```
  
      Args:
        x: A rank 1 `Tensor` of any dtype for which arithmetic operations are
          permitted.
        segment_ids: A `Tensor`. Must be one of the following types: int32, int64. A
          1-D tensor whose size is equal to the size of `x`. Values should be sorted
          and can be repeated. Values must range from `0` to `num segments - 1`.
        exclusive: Python bool. See description above.
          Default value: False
        dtype: Optional `tf.Dtype`. If supplied, the dtype for `x` to use when
          converting to `Tensor`.
          Default value: None which maps to the default dtype inferred by TF.
        name: Python `str` name prefixed to Ops created by this class.
          Default value: None which is mapped to the default name 'segment_cumsum'.
  
      Returns:
        cumsums: A `Tensor` of the same dtype as `x`. Assuming that each segment is
          of length greater than or equal to order, if `exclusive` is True,
          then the size is `n-order*k` where `n` is the size of x,
          `k` is the number of different segment ids supplied if `segment_ids` is
          not None or 1 if `segment_ids` is None. If any of the segments is of
          length less than the order, then the size is:
          `n-sum(min(order, length(segment_j)), j)` where the sum is over segments.
          If `exclusive` is False, then the size is `n`.
      """
      with tf.compat.v1.name_scope(
          name, default_name="segment_cumsum", values=[x]
      ):
          x = tf.convert_to_tensor(x, dtype=dtype)
          raw_cumsum = tf.math.cumsum(x, exclusive=exclusive)
          if segment_ids is None:
              return raw_cumsum
  
          # It is quite tedious to do a vectorized version without a while loop so
          # we skip that for now.
          # TODO(b/137940928): Replace these ops with more efficient C++ kernels.
          def scanner(accumulators, args):
              cumsum, prev_segment, prev_value = accumulators
              value, segment = args
              if exclusive:
                  initial_value, inc_value = (
                      tf.zeros_like(value),
                      cumsum + prev_value,
                  )
              else:
                  initial_value, inc_value = value, cumsum + value
              next_cumsum = tf.where(
                  tf.equal(prev_segment, segment),
                  inc_value,
                  initial_value,
              )
              return next_cumsum, segment, value
  
          return tf.scan(
              scanner,
              (x, segment_ids),
              initializer=(
                  tf.zeros_like(x[0]),
                  tf.zeros_like(segment_ids[0]) - 1,
                  tf.zeros_like(x[0]),
              ),
          )[0]
  
imports: import tensorflow.compat.v2
  from tf_quant_finance.math import diff_ops
label_desc: Use image node_id 2 for calling a global function with example usage: segment_cumsum(x, segment_ids, exclusive, dtype, name) and returns: tf, raw_cumsum, next_cumsum,  segment,  value