def cumsum(x, dim, exclusive=False):
  """Cumulative sum.

  Args:
    x: a Tensor
    dim: a Dimension
    exclusive: a boolean

  Returns:
    a Tensor with the same shape as x.
  """
  with tf.variable_scope("cumsum"):
    new_name = "tmp_dim_cumsum"
    new_dim = Dimension(new_name, dim.size)
    new_shape = x.shape.rename_dimension(dim.name, new_name)
    comparator = less if exclusive else less_equal
    m = cast(
        comparator(mtf_range(x.mesh, dim, dtype=tf.float32),
                   mtf_range(x.mesh, new_dim, dtype=tf.float32)), x.dtype)
    ret = einsum([x, m], output_shape=new_shape)
    return reshape(ret, x.shape)
Cumulative sum.

  Args:
    x: a Tensor
    dim: a Dimension
    exclusive: a boolean

  Returns:
    a Tensor with the same shape as x.