code: def cumsum(x, axis=0, exclusive=False):
    
    if not is_xla_compiled():
      return tf.cumsum(x, axis=axis, exclusive=exclusive)
    x_shape = shape_list(x)
    rank = len(x_shape)
    length = x_shape[axis]
    my_range = tf.range(length)
    comparator = tf.less if exclusive else tf.less_equal
    mask = tf.cast(
        comparator(tf.expand_dims(my_range, 1), tf.expand_dims(my_range, 0)),
        x.dtype)
    ret = tf.tensordot(x, mask, axes=[[axis], [0]])
    if axis != rank - 1:
      ret = tf.transpose(
          ret,
          list(range(axis)) + [rank - 1] + list(range(axis, rank - 1)))
    return ret
docs: TPU hack for tf.cumsum.
  
    This is equivalent to tf.cumsum and is faster on TPU as of 04/2018 unless
    the axis dimension is very large.
  
    Args:
      x: a Tensor
      axis: an integer
      exclusive: a boolean
  
    Returns:
      Tensor of the same shape as x.