def _sparse_reduce_batch_keep_shape(
    sparse_reduce_fn: Callable[..., tf.Tensor], sparse_tensor: tf.SparseTensor
) -> tf.Tensor:  # pylint: disable=g-bare-generic
  """Applies a tf.sparse.reduce_* method on the given sparse_tensor."""
  result = sparse_reduce_fn(sparse_tensor, axis=0)
  result.set_shape(sparse_tensor.get_shape()[1:])
  return result
Applies a tf.sparse.reduce_* method on the given sparse_tensor.
tensorflow_transform/tf_utils.py