def aggregate_sparse_indices(indices, values, shape, agg_fn="sum"):
  """Sums values corresponding to repeated indices.

  Returns the unique indices and their summed values.

  Args:
    indices: [num_nnz, rank] Tensor.
    values: [num_nnz] Tensor.
    shape: [rank] Tensor.
    agg_fn: Method to use for aggregation - `sum` or `max`.

  Returns:
    indices: [num_uniq, rank] Tensor.
    values: [num_uniq] Tensor.
  """
  # Linearize the indices.
  scaling_vec = tf.cumprod(tf.cast(shape, indices.dtype), exclusive=True)
  linearized = tf.linalg.matvec(indices, scaling_vec)
  # Get the unique indices, and their positions in the array
  y, idx = tf.unique(linearized)
  # Use the positions of the unique values as the segment ids to
  # get the unique values
  idx.set_shape([None])
  if agg_fn == "sum":
    values = tf.unsorted_segment_sum(values, idx, tf.shape(y)[0])
  elif agg_fn == "max":
    values = tf.unsorted_segment_max(values, idx, tf.shape(y)[0])
  elif agg_fn == "mean":
    values = tf.unsorted_segment_mean(values, idx, tf.shape(y)[0])
  # Go back to ND indices
  y = tf.expand_dims(y, 1)
  indices = tf.floormod(
      tf.floordiv(y, tf.expand_dims(scaling_vec, 0)),
      tf.cast(tf.expand_dims(shape, 0), indices.dtype))
  return indices, values
Sums values corresponding to repeated indices.

  Returns the unique indices and their summed values.

  Args:
    indices: [num_nnz, rank] Tensor.
    values: [num_nnz] Tensor.
    shape: [rank] Tensor.
    agg_fn: Method to use for aggregation - `sum` or `max`.

  Returns:
    indices: [num_uniq, rank] Tensor.
    values: [num_uniq] Tensor.
  
language/labs/drkit/model_fns.py