def sparse_reduce(sp_tensor, rank, agg_fn="sum", axis=-1):
  """Reduce SparseTensor along the given axis.

  Args:
    sp_tensor: SparseTensor of arbitrary rank.
    rank: Integer rank of the sparse tensor.
    agg_fn: Reduce function for aggregation.
    axis: Integer specifying axis to sum over.

  Returns:
    sp_tensor: SparseTensor of one less rank.
  """
  if axis < 0:
    axis = rank + axis
  axes_to_keep = tf.one_hot(
      axis, rank, on_value=False, off_value=True, dtype=tf.bool)
  indices_to_keep = tf.boolean_mask(sp_tensor.indices, axes_to_keep, axis=1)
  new_shape = tf.boolean_mask(sp_tensor.dense_shape, axes_to_keep)
  indices_to_keep.set_shape([None, rank - 1])
  indices, values = aggregate_sparse_indices(
      indices_to_keep, sp_tensor.values, new_shape, agg_fn=agg_fn)
  return tf.sparse.reorder(tf.SparseTensor(indices, values, new_shape))
Reduce SparseTensor along the given axis.

  Args:
    sp_tensor: SparseTensor of arbitrary rank.
    rank: Integer rank of the sparse tensor.
    agg_fn: Reduce function for aggregation.
    axis: Integer specifying axis to sum over.

  Returns:
    sp_tensor: SparseTensor of one less rank.
  
language/labs/drkit/model_fns.py