// CHECK-DAG: %[[cst_1:.*]] = arith.constant dense<[-1, 36, 32]> : tensor<3xi64>
// CHECK: %[[v0:.*]] = "tf.Shape"(%arg0) : (tensor<?x36x?xf32>) -> tensor<3xi32>
// CHECK: %[[v1:.*]] = "tf.UnsortedSegmentProd"(%0, %[[cst]], %cst_0) : (tensor<3xi32>, tensor<3xi32>, tensor<i32>) -> tensor<4xi32>
// CHECK: %[[v2:.*]] = "tf.Reshape"(%arg0, %1) : (tensor<?x36x?xf32>, tensor<4xi32>) -> tensor<?x36x1x?xf32>

}

// Creates ReshapeOp with runtime calcuation of required shape to support
// dynamic shapes. The shape is calculated by Shape and UnsortedSegmentProd op.
// `reshape_segids` and `num_reshape_segids` for UnsortedSegmentProd is
// calculated in `reshapeForBatchMatmul`.
TF::ReshapeOp createReshapeOpForDynamic(Value value, ArrayRef<int64_t> shape,

// CHECK-DAG:       %cst_2 = "tf.Const"() <{value = dense<1> : tensor<i32>}> : () -> tensor<i32>
// CHECK:           %2 = "tf.UnsortedSegmentProd"(%1, %cst_0, %cst_2) : (tensor<3xi32>, tensor<3xi32>, tensor<i32>) -> tensor<1xi32>
// CHECK:           %3 = "tf.UnsortedSegmentProd"(%1, %cst_1, %cst_2) : (tensor<3xi32>, tensor<3xi32>, tensor<i32>) -> tensor<1xi32>

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

}

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

}

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

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

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

  // expected-error @+1 {{requires segment ids shape to be a prefix of data shape, but dimension #1 differs: 9 vs. 10}}
  %0 = "tf.UnsortedSegmentProd"(%data, %segment_ids, %num_segments) : (tensor<7x10x8xf32>, tensor<7x9xi32>, tensor<i32>) -> (tensor<?x8xf32>)
  func.return
}

// -----

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

// -----

// CHECK-LABEL: @unsorted_segment_min

  TF_DerivedOperandTypeAttr Tnumsegments = TF_DerivedOperandTypeAttr<2>;

  let hasVerifier = 1;
}

def TF_UnsortedSegmentProdOp : TF_Op<"UnsortedSegmentProd", [Pure]> {
  let summary = "Computes the product along segments of a tensor.";

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

*   `tf.lite`:

    *   New operations supported:
        *   tflite SelectV2 now supports 5D.
        *   `tf.einsum` is supported with multiple unknown shapes.
        *   `tf.unsortedsegmentprod` op is supported.
        *   `tf.unsortedsegmentmax` op is supported.
        *   `tf.unsortedsegmentsum` op is supported.
    *   Updates to existing operations: