# RUN: not tf_tfl_translate -tf-input-arrays=input0 -tf-input-shapes=-1 -tf-input-data-types=DT_FLOAT -tf-output-arrays=add %s 2>&1 | FileCheck %s
# CHECK: error: The Graph contains unsupported `StatefulPartionedCallOp`(s)

node {
  name: "input0"
  op: "Placeholder"
  attr {
    key: "dtype"
    value {
      type: DT_FLOAT
    }
  }
}
node {
  name: "args_0"
  op: "_Arg"
  attr {
    key: "T"
    value {

# RUN: tf_tfl_translate -tf-input-arrays=input0,input1 \
# RUN:                  -tf-input-shapes=4:4 \
# RUN:                  -tf-input-data-types=DT_FLOAT,DT_FLOAT \
# RUN:                  -tf-output-arrays=Add \
# RUN:                  -tf-inference-type=DT_QUINT8 \
# RUN:                  -tf-input-min-values='-2,-3' \
# RUN:                  -tf-input-max-values='2,3' \
# RUN:                  --quant-stats=%s.stats \

# RUN: tf_tfl_translate -tf-input-arrays=input0,input1 -tf-input-shapes=4:4 -tf-input-data-types=DT_INT32,DT_INT32 -tf-output-arrays=Add %s -o - | flatbuffer_to_string - | FileCheck %s

# Add two tensor<4xi32> inputs and return the result

node {
  name: "Add"
  op: "Add"
  input: "input0"
  input: "input1"
  attr {
    key: "T"
    value {
      type: DT_INT32
    }
  }
}
node {
  name: "input0"
  op: "Placeholder"
  attr {

}

// CHECK:      node {
// CHECK-NEXT:   name: "input0"
// CHECK-NEXT:   op: "_Arg"
// CHECK:      node {
// CHECK-NEXT:   name: "input1"
// CHECK-NEXT:   op: "_Arg"
// CHECK:      node {
// CHECK-NEXT:   name: "Add{{_.*_1}}"
// CHECK-NEXT:   op: "Add"
// CHECK-NEXT:   input: "input0"
// CHECK-NEXT:   input: "input1"
// CHECK:      node {
// CHECK-NEXT:   name: "Add"
// CHECK-NEXT:   op: "_Retval"

  %10 = "tf.opK"() : () -> tensor<*xi16>
  %11 = "tf.opL"() : () -> tensor<*xi64>
  tf_device.replicate([%0, %1, %2] as %input0: tensor<*xi1>, %9 as %input1: tensor<*xi8>, %10 as %input2: tensor<*xi16>, [%3, %4, %5] as %input3: tensor<*xi32>, [%6, %7, %8] as %input4: tensor<*xf32>, %11 as %input5: tensor<*xi64>) {n = 3 : i32} {
    tf_device.return
  }
  func.return

// CHECK:      %[[OP_A:[a-z0-9]*]] = "tf.opA"

  }
}
node {
  name: "ParseExample/ParseExampleV2"
  op: "ParseExampleV2"
  input: "input0"
  input: "ParseExample/ParseExampleV2/names"
  input: "ParseExample/ParseExampleV2/sparse_keys"
  input: "ParseExample/ParseExampleV2/dense_keys"
  input: "ParseExample/ParseExampleV2/ragged_keys"
  input: "ParseExample/Const"
  input: "ParseExample/Const_1"
  attr {
    key: "Tdense"
    value {
      list {

>;

// Removes `stablehlo.custom_call @Sharding`. Assumes this call always accepts
// one input. Other attributes are ignored.
def RemoveShardingCustomCall : Pat<
  (StableHLO_CustomCallOp (variadic $input0), $call_target_name, $_, $_, $_, $_, $_, $_, $_),
  (replaceWithValue $input0),

// RUN: tac-translate -input-mlir -output-mlir -device-specs=GPU %s -o - 2>&1 | FileCheck %s

module {
func.func @main(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>, %arg2: tensor<1xf32>, %arg3: tensor<1xf32>) -> tensor<2x1xf32> attributes {tf.entry_function = {inputs = "input0,input1,input2,input3", outputs = "output"}} {
  %0 = "tfl.add"(%arg0, %arg1) {fused_activation_function = "RELU6"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>

      // CHECK-NEXT: op: "ParseExampleV2"
      // CHECK-NEXT: input: "input0"
      // CHECK-NEXT: input: "tf.Const{{_.*_3}}"
      // CHECK-NEXT: input: "tf.Const{{_.*_5}}"
      // CHECK-NEXT: input: "tf.Const{{_.*_2}}"
      // CHECK-NEXT: input: "tf.Const{{_.*_4}}"
      // CHECK-NEXT: input: "tf.Const"
      // CHECK-NEXT: input: "tf.Const{{_.*_1}}"
      // CHECK-NEXT: attr {
      // CHECK-NEXT:   key: "Tdense"

      %outputs_5, %control_6 = tf_executor.island wraps "tf.TPUReplicatedInput"(%outputs_3) {N = 1 : i64, T = "tfdtype$DT_FLOAT", device = "", name = "input0"} : (tensor<0xf32>) -> tensor<0xf32>