// RUN: flatbuffer_translate -mlir-to-tflite-flatbuffer %s -o - | flatbuffer_to_string - | FileCheck %s

func.func @main(tensor<8xi32>, tensor<8xi32>, tensor<i32>) -> tensor<8xi32> {
^bb0(%arg0: tensor<8xi32>, %arg1: tensor<8xi32>, %arg2: tensor<i32>):
// CHECK: {
// CHECK-NEXT:   version: 3,
// CHECK-NEXT:   operator_codes: [ {
// CHECK-NEXT:     deprecated_builtin_code: 127,
// CHECK-NEXT:     version: 1,
// CHECK-NEXT:     builtin_code: UNSORTED_SEGMENT_PROD

    ```mlir
    func @tf_tpu_rewrite(%arg0: tensor<8xi32>) -> tensor<8xi32> {
      %0 = "tf_device.cluster_func"(%arg0) {_xla_compile_device_type = "TPU", _replication_info = "cluster0", func = @func, num_cores_per_replica = 2, input_sharding_configuration = ["\08\01\1A\01\01\22\01\00"], output_sharding_configuration = ["\08\01\1A\01\01\22\01\00"]} : (tensor<8xi32>) -> tensor<8xi32>
      return %0 : tensor<8xi32>
    }
    ```

    will be rewritten as:

  attr {
    key: "dtype"
    value {
      type: DT_INT32
    }
  }
}
versions {
  producer: 27
}

# CHECK-LABEL: func @main
# CHECK-SAME:  (%[[ARG_0:[a-z0-9]+]]: tensor<8xi32>) -> tensor<8xi32>
# CHECK-SAME:  control_outputs = ""
# CHECK-SAME:  inputs = "input"
# CHECK-SAME:  outputs = "input"
# CHECK:         %[[GRAPH:[0-9]+]] = tf_executor.graph
# CHECK:           tf_executor.fetch %[[ARG_0]]

    }
  }
  attr {
    key: "index"
    value {
      i: 0
    }
  }
}
versions {
  producer: 27
}

# CHECK-LABEL: func @main
# CHECK-SAME:  (%[[ARG_0:[a-z0-9]+]]: tensor<8xi32> loc({{.*}})) -> tensor<8xi32>
# CHECK-SAME:  control_outputs = ""
# CHECK-SAME:  inputs = "arg"
# CHECK-SAME:  outputs = "arg"
# CHECK:         %[[GRAPH:[0-9]+]] = tf_executor.graph
# CHECK:           tf_executor.fetch %[[ARG_0]]

  func.return %0#0 : tensor<2xi32>

// CHECK-LABEL: unpack3Tensors
// CHECK: "tfl.unpack"(%arg0) <{axis = 1 : i32, num = 3 : i32}> : (tensor<2x3xi32>) -> (tensor<2xi32>, tensor<2xi32>, tensor<2xi32>)
}

func.func @unpackNegAxis(%arg0: tensor<2x3xi32>) -> tensor<2xi32> {
  %0:3 = "tf.Unpack"(%arg0) {axis = -1 : i64} : (tensor<2x3xi32>) -> (tensor<2xi32>, tensor<2xi32>, tensor<2xi32>)

    // CHECK:      "tf._TPUCompileMlir"(%[[ARG_0_SHAPE]], %[[ARG_2_SHAPE]])

    func.return %0: tensor<8xi32>
  }
  func.func @_func(%arg0: tensor<*xi32>, %arg1: tensor<8xi32>, %arg2: tensor<*xi32>, %arg3: tensor<8xi32>) -> tensor<8xi32> {
    func.return %arg1 : tensor<8xi32>
  }
}

// -----

    >} : (
      tensor<?x?x6xi8>, tensor<6x8xi8>
    ) -> tensor<?x?x8xi32>
  %1 = chlo.broadcast_add %0, %zp_offset : (
      tensor<?x?x8xi32>, tensor<8xi32>) -> tensor<?x?x8xi32>
  %2 = chlo.broadcast_add %1, %bias : (
      tensor<?x?x8xi32>, tensor<8xi32>) -> tensor<?x?x8xi32>
  return %2 : tensor<?x?x8xi32>
}

// -----

// CHECK-LABEL: func @dot_general_add_add_static

  func.return %0#0 : tensor<2xi32>
}

// -----

func.func @unpack(%arg0: tensor<*xi32>) -> (tensor<*xi32>, tensor<*xi32>) {
  %0:2 = "tfl.unpack"(%arg0) {axis = 1 : i32, num = 2 : i32} : (tensor<*xi32>) -> (tensor<*xi32>, tensor<*xi32>)
  func.return %0#0, %0#1 : tensor<*xi32>, tensor<*xi32>
}

// -----

// CHECK-NEXT:      %6 = "tfl.pseudo_const"() <{value = dense<1> : tensor<i32>}> : () -> tensor<i32>
// CHECK-NEXT:      %7 = "tfl.unsorted_segment_prod"(%3, %4, %6) : (tensor<4xi32>, tensor<4xi32>, tensor<i32>) -> tensor<1xi32>
// CHECK-NEXT:      %8 = "tfl.unsorted_segment_prod"(%3, %5, %6) : (tensor<4xi32>, tensor<4xi32>, tensor<i32>) -> tensor<1xi32>

  func.return %0#0 : tensor<8x8x8x8xf32>
}

// -----

// Test invalid tf.FusedBatchNorm
func.func @testFusedBatchNormWrongXType(tensor<8x8x8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>) -> tensor<8x8x8xf32> {