}

// -----

func.func @unsorted_segment_sum(%arg0: tensor<8xf32>, %arg1: tensor<8xi32>) -> tensor<8xf32> {
  %num_segments = "tf.Const"() {value = dense<8> : tensor<i32>} : () -> tensor<i32>
  %0 = "tf.UnsortedSegmentSum"(%arg0, %arg1, %num_segments) : (tensor<8xf32>, tensor<8xi32>, tensor<i32>) -> tensor<8xf32>
  func.return %0 : tensor<8xf32>
  // CHECK-LABEL: unsorted_segment_sum

  BUCKETIZE = 147,
  RANDOM_UNIFORM = 148,
  MULTINOMIAL = 149,
  GELU = 150,
  DYNAMIC_UPDATE_SLICE = 151,
  RELU_0_TO_1 = 152,
  UNSORTED_SEGMENT_PROD = 153,
  UNSORTED_SEGMENT_MAX = 154,
  UNSORTED_SEGMENT_SUM = 155,
  ATAN2 = 156,
  UNSORTED_SEGMENT_MIN = 157,
  SIGN = 158,
  BITCAST = 159,
  BITWISE_XOR = 160,
  RIGHT_SHIFT = 161,
  // All Operators start with STABLEHLO_ prefixes are subject to change

// -----

// CHECK-LABEL: testUnsortedSegmentSum
func.func @testUnsortedSegmentSum(%arg0: tensor<8xf32>, %arg1: tensor<8xi32>,  %arg2: tensor<i32>) -> tensor<8xf32> {
  // CHECK: "tfl.unsorted_segment_sum"(%arg0, %arg1, %arg2)
  %0 = "tfl.unsorted_segment_sum"(%arg0, %arg1, %arg2) : (tensor<8xf32>, tensor<8xi32>, tensor<i32>) -> tensor<8xf32>
  func.return %0 : tensor<8xf32>
  // CHECK: return %0 : tensor<8xf32>
}

// -----

//===----------------------------------------------------------------------===//

// -----

// CHECK-LABEL: @unsorted_segment_sum
// CHECK-SAME: [[DATA:%.*]]: tensor<8x16x64xf32>
// CHECK-SAME: [[SI:%.*]]: tensor<8x16xi32>
func.func @unsorted_segment_sum(%data: tensor<8x16x64xf32>, %segment_ids : tensor<8x16xi32>) -> (tensor<4x64xf32>) {

  );

  let results = (outs TFL_TensorOf<[F32, I32]>:$output);
}

def TFL_UnsortedSegmentSumOp: TFL_Op<"unsorted_segment_sum", [
    Pure, PredOpTrait<"input and output must have same element type",
      TFL_TCresVTEtIsSameAsOp<0, 0>>]> {

  let summary = "UnsortedSegmentSum operator";

  let description = [{
    From a tensor segmentation, computes the `output` resulting from

cluster 29 size 9
 Add 1
 Mul 2
 RealDiv 2
 Sqrt 2
 Sub 2
cluster 30 size 9
 Add 1
 Mul 2
 RealDiv 2
 Sqrt 2
 Sub 2
cluster 31 size 4
 Mul 3
 UnsortedSegmentSum 1
cluster 32 size 4
 Mul 3

        *   (since v2.7) Stateful ops used in `tf.data.Dataset`
        *   (since v2.7) `tf.convert_to_tensor` when fed with (sparse)
            `tf.IndexedSlices` (because it uses `tf.math.unsorted_segment_sum`)
        *   (since v2.7) `tf.gather` backprop (because `tf.convert_to_tensor`
            reduces `tf.gather`'s (sparse) `tf.IndexedSlices` gradients into its
            dense `params` input)

<div style="width:70%; margin:auto; margin-bottom:10px; margin-top:20px;">
<img style="width:100%" src="https://www.tensorflow.org/images/UnsortedSegmentSum.png" alt>
</div>

>>> c = [[1,2,3,4], [5,6,7,8], [4,3,2,1]]
>>> tf.math.unsorted_segment_sum(c, [0, 1, 0], num_segments=2).numpy()
array([[5, 5, 5, 5],
       [5, 6, 7, 8]], dtype=int32)
  }];

  let arguments = (ins

                              std::vector<Output>* grad_outputs) {
  if (op.num_inputs() != 3) {
    return errors::InvalidArgument("UnsortedSegmentSum requires 3 arguments");
  }

  if (grad_inputs.size() != 1) {
    return errors::InvalidArgument(
        "UnsortedSegmentSum grad requires 1 grad input");
  }

  auto zero_clipped_indices = GetZeroClippedIndices(scope, op.input(1));

using ops::RealDiv;
using ops::SegmentSum;
using ops::SelectV2;
using ops::SquaredDifference;
using ops::Sub;
using ops::Sum;
using ops::UnsortedSegmentMax;
using ops::UnsortedSegmentMin;
using ops::UnsortedSegmentSum;
using ops::Where3;

// TODO(andydavis) Test gradient function against numeric gradients output.
// TODO(andydavis) As more gradients are added move common test functions
// to a testutil library.