Add back list tf.sparse.segment_sum docstring, and fix horrible bug i…

…n it which would compute segment_mean when segment_sum was called.

PiperOrigin-RevId: 250559891

tensorflow/python/ops/math_ops.py

@@ -3416,8 +3416,8 @@ def sparse_segment_sum(data, indices, segment_ids, name=None,
   segmentation](https://tensorflow.org/api_docs/python/tf/math#Segmentation)
   for an explanation of segments.
 
-  Like `SegmentSum`, but `segment_ids` can have rank less than `data`'s first
-  dimension, selecting a subset of dimension 0, specified by `indices`.
+  Like `tf.math.segment_sum`, but `segment_ids` can have rank less than `data`'s
+  first dimension, selecting a subset of dimension 0, specified by `indices`.
   `segment_ids` is allowed to have missing ids, in which case the output will
   be zeros at those indices. In those cases `num_segments` is used to determine
   the size of the output.
@@ -3486,7 +3486,65 @@ def sparse_segment_sum_v2(data,
                           segment_ids,
                           num_segments=None,
                           name=None):
-  return sparse_segment_mean(
+  r"""Computes the sum along sparse segments of a tensor.
+
+  Read [the section on
+  segmentation](https://tensorflow.org/api_docs/python/tf/math#Segmentation)
+  for an explanation of segments.
+
+  Like `tf.math.segment_sum`, but `segment_ids` can have rank less than `data`'s
+  first dimension, selecting a subset of dimension 0, specified by `indices`.
+  `segment_ids` is allowed to have missing ids, in which case the output will
+  be zeros at those indices. In those cases `num_segments` is used to determine
+  the size of the output.
+
+  For example:
+
+  ```python
+  c = tf.constant([[1,2,3,4], [-1,-2,-3,-4], [5,6,7,8]])
+
+  # Select two rows, one segment.
+  tf.sparse.segment_sum(c, tf.constant([0, 1]), tf.constant([0, 0]))
+  # => [[0 0 0 0]]
+
+  # Select two rows, two segment.
+  tf.sparse.segment_sum(c, tf.constant([0, 1]), tf.constant([0, 1]))
+  # => [[ 1  2  3  4]
+  #     [-1 -2 -3 -4]]
+
+  # With missing segment ids.
+  tf.sparse.segment_sum(c, tf.constant([0, 1]), tf.constant([0, 2]),
+                        num_segments=4)
+  # => [[ 1  2  3  4]
+  #     [ 0  0  0  0]
+  #     [-1 -2 -3 -4]
+  #     [ 0  0  0  0]]
+
+  # Select all rows, two segments.
+  tf.sparse.segment_sum(c, tf.constant([0, 1, 2]), tf.constant([0, 0, 1]))
+  # => [[0 0 0 0]
+  #     [5 6 7 8]]
+
+  # Which is equivalent to:
+  tf.math.segment_sum(c, tf.constant([0, 0, 1]))
+  ```
+
+  Args:
+    data: A `Tensor` with data that will be assembled in the output.
+    indices: A 1-D `Tensor` with indices into `data`. Has same rank as
+      `segment_ids`.
+    segment_ids: A 1-D `Tensor` with indices into the output `Tensor`. Values
+      should be sorted and can be repeated.
+    num_segments: An optional int32 scalar. Indicates the size of the output
+      `Tensor`.
+    name: A name for the operation (optional).
+
+  Returns:
+    A `tensor` of the shape as data, except for dimension 0 which
+    has size `k`, the number of segments specified via `num_segments` or
+    inferred for the last element in `segments_ids`.
+  """
+  return sparse_segment_sum(
       data, indices, segment_ids, name=name, num_segments=num_segments)
 
 
@@ -3503,8 +3561,9 @@ def sparse_segment_mean(data,
   segmentation](https://tensorflow.org/api_docs/python/tf/math#Segmentation)
   for an explanation of segments.
 
-  Like `SegmentMean`, but `segment_ids` can have rank less than `data`'s first
-  dimension, selecting a subset of dimension 0, specified by `indices`.
+  Like `tf.math.segment_mean`, but `segment_ids` can have rank less than
+  `data`'s first dimension, selecting a subset of dimension 0, specified by
+  `indices`.
   `segment_ids` is allowed to have missing ids, in which case the output will
   be zeros at those indices. In those cases `num_segments` is used to determine
   the size of the output.
@@ -3548,8 +3607,9 @@ def sparse_segment_mean_v2(data,
   segmentation](https://tensorflow.org/api_docs/python/tf/math#Segmentation)
   for an explanation of segments.
 
-  Like `SegmentMean`, but `segment_ids` can have rank less than `data`'s first
-  dimension, selecting a subset of dimension 0, specified by `indices`.
+  Like `tf.math.segment_mean`, but `segment_ids` can have rank less than
+  `data`'s first dimension, selecting a subset of dimension 0, specified by
+  `indices`.
   `segment_ids` is allowed to have missing ids, in which case the output will
   be zeros at those indices. In those cases `num_segments` is used to determine
   the size of the output.
@@ -3619,7 +3679,12 @@ def sparse_segment_sqrt_n_v2(data,
                              name=None):
   r"""Computes the sum along sparse segments of a tensor divided by the sqrt(N).
 
-  `N` is the size of the segment being reduced.
+  Read [the section on
+  segmentation](https://tensorflow.org/api_docs/python/tf/math#Segmentation)
+  for an explanation of segments.
+
+  Like `tf.sparse.segment_mean`, but instead of dividing by the size of the
+  segment, `N`, divide by `sqrt(N)` instead.
 
   Args:
     data: A `Tensor` with data that will be assembled in the output.