}


// -----

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

  RANDOM_STANDARD_NORMAL = 146,
  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,

// -----

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

// -----

  func.return %0: tensor<4x?xf32>
}

// -----

// CHECK-LABEL: @unsorted_segment_max
func.func @unsorted_segment_max(%data: tensor<8x?x64xf32>, %segment_ids : tensor<?x16xi32>) -> (tensor<4x?xf32>) {
  %num_segments = "tf.Const"() {value = dense<4> : tensor<i32>} : () -> tensor<i32>

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

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

  let summary = "UnsortedSegmentMax operator";

  let description = [{

    "RANDOM_STANDARD_NORMAL",
    "BUCKETIZE",
    "RANDOM_UNIFORM",
    "MULTINOMIAL",
    "GELU",
    "DYNAMIC_UPDATE_SLICE",
    "RELU_0_TO_1",
    "UNSORTED_SEGMENT_PROD",
    "UNSORTED_SEGMENT_MAX",
    "UNSORTED_SEGMENT_SUM",
    "ATAN2",
    "UNSORTED_SEGMENT_MIN",
    "SIGN",
    "BITCAST",
    "BITWISE_XOR",
    "RIGHT_SHIFT",
    "STABLEHLO_LOGISTIC",
    "STABLEHLO_ADD",

<img style="width:100%" src="https://www.tensorflow.org/images/UnsortedSegmentMax.png" alt>
</div>

For example:

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

  let arguments = (ins
    TF_IntOrFpTensor:$data,

        `tf.math.sin`, `tf.math.sinh`, `tf.math.softplus`, `tf.math.softsign`,
        `tf.math.squared_difference`, `tf.math.tan`,
        `tf.math.unsorted_segment_max`, `tf.math.unsorted_segment_min`,
        `tf.math.unsorted_segment_prod`, `tf.math.unsorted_segment_sum`,
        `tf.math.zeta`.
    *   New endpoints in `tf.quantization` namespace: