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

                              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.

    values = Reshape(scope, values, flat_values_shape);
  }

  indices = Reshape(scope, indices, indices_size);
  Output params_grad =
      UnsortedSegmentSum(scope, values, indices, gather_dim_size);

  if (batch_dims != 0) {
    // Put back the batch dimensions.
    params_grad = Reshape(scope, params_grad, params_shape);
  }
  return params_grad;
}

}

// -----

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

      "UniformRequantize",
      "Unique",
      "UniqueV2",
      "UpperBound",
      "UnsortedSegmentMax",
      "UnsortedSegmentMin",
      "UnsortedSegmentProd",
      "UnsortedSegmentSum",
      "VarIsInitializedOp",
      "VariableShape",
      "Where",
      "While",
      "XlaAllReduce",
      "XlaBroadcastHelper",
      "XlaCallModule",
      "XlaConcatND",

func.func @unsortedSegmentReduction(%data: tensor<7x10x8xf32>, %segment_ids: tensor<7x10x8x1xi32>, %num_segments: tensor<i32>) {
  // expected-error @+1 {{requires segment ids rank to be less than or equal to data's rank}}
  %0 = "tf.UnsortedSegmentSum"(%data, %segment_ids, %num_segments) : (tensor<7x10x8xf32>, tensor<7x10x8x1xi32>, tensor<i32>) -> (tensor<?x8xf32>)
  func.return
}

// -----

}

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
    summing together elements mapped to the same segment_id. I.e. `output[i]` is

  // CHECK:   mhlo.return [[ADD]]
  // CHECK-NEXT: (tensor<4x64xf32>, tensor<8x16xi32>, tensor<8x16x64xf32>) -> tensor<4x64xf32>
  // CHECK: return [[SCATTER]]
  %0 = "tf.UnsortedSegmentSum"(%data, %segment_ids, %num_segments) : (tensor<8x16x64xf32>, tensor<8x16xi32>, tensor<i32>) -> (tensor<4x64xf32>)
  func.return %0: tensor<4x64xf32>
}

// -----

// CHECK-LABEL: @unsorted_segment_prod

  TF_DerivedOperandTypeAttr Tnumsegments = TF_DerivedOperandTypeAttr<2>;

  let hasVerifier = 1;
}

def TF_UnsortedSegmentSumOp : TF_Op<"UnsortedSegmentSum", [Pure]> {
  let summary = "Computes the sum along segments of a tensor.";

  let description = [{
Read
[the section on segmentation](https://tensorflow.org/api_docs/python/tf/math#Segmentation)