TensorShape y_shape({2, 1, 2});
  RunTest({x}, {x_shape}, {y}, {y_shape});
}

TEST_F(NaryGradTest, SegmentSum) {
  TensorShape x_shape({3, 4});
  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
  auto y = SegmentSum(scope_, x, {0, 0, 1});
  // the sum is always on the first dimension
  TensorShape y_shape({2, 4});
  RunTest({x}, {x_shape}, {y}, {y_shape});
}

  // The SegmentSum operation sums segments of the Tensor that have the same
  // index in the segment_ids parameter.
  // i.e z = [2, 3, 4, 5], segment_ids [0, 0, 0, 1]
  // will produce [2 + 3 + 4, 5] = [9, 5]
  // The gradient that will flow back to the gather operation will look like
  // [x1, x2], it will have the same shape as the output of the SegmentSum

  // CHECK-LABEL: cumsum_i64
  // CHECK: "tfl.cast"
  // CHECK: "tfl.cumsum"
}

func.func @segmentsum(%arg0: tensor<3x3xf32>, %arg1: tensor<i32>) -> tensor<*xf32> {
  %0 = "tf.SegmentSum"(%arg0, %arg1) : (tensor<3x3xf32>, tensor<i32>) -> tensor<*xf32>
  func.return %0 : tensor<*xf32>
  // CHECK-LABEL: segmentsum
  // CHECK: "tfl.segment_sum"(%arg0, %arg1) : (tensor<3x3xf32>, tensor<i32>) -> tensor<*xf32>
}

    }
  }];
}

def TFL_SegmentSumOp: TFL_Op<"segment_sum", [
    Pure,
    PredOpTrait<"input and output must have same element type",
      TFL_TCresVTEtIsSameAsOp<0, 0>>]> {
  let summary = "SegmentSum operator";

  let description = [{
    Computes the sum along segments of a tensor.
  }];

  let arguments = (ins
    TFL_TensorOf<[F32, I32]>:$input,
    TFL_I32Tensor:$segment_ids
  );

  TF_DerivedOperandTypeAttr Tindices = TF_DerivedOperandTypeAttr<1>;
  TF_DerivedOperandTypeAttr Tnumsegments = TF_DerivedOperandTypeAttr<2>;
}

def TF_SegmentSumOp : TF_Op<"SegmentSum", [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)