}

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,

  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: