// Tests input and output names from FlatBuffer are added to `tf.entry_function` attribute.

// CHECK-LABEL: @main
func.func @main(%arg0: tensor<4xi8>, %arg1: tensor<4xi32>) -> (tensor<4xi32>, tensor<4xi8>)
// CHECK: attributes {tf.entry_function = {inputs = "input0,input1", outputs = "output0,output1"}}
attributes {tf.entry_function = {inputs = "input0,input1", outputs = "output0,output1"}} {

          ],
          "name": "output0",
          "quantization": {
          }
        },
        {
          "shape": [
            128,
            32,
            32,
            3
          ],
          "name": "output1",
          "quantization": {
          }
        }
      ],
      "inputs": [
        0,
        1
      ],
      "outputs": [
        2,
        3
      ],

// RUN: tf-mlir-translate -mlir-to-graphdef %s -o - | FileCheck %s

func.func @main(%arg0: tensor<10xi32>) -> tensor<10xi32>
attributes {tf.entry_function = {inputs = "input0", outputs = "output0"}} {
  %graph = tf_executor.graph {
    tf_executor.fetch %arg0 : tensor<10xi32>
  }
  func.return %graph : tensor<10xi32>
}

// CHECK:      node {
// CHECK-NEXT:   name: "input0"
// CHECK-NEXT:   op: "_Arg"
// CHECK:          key: "T"

  %graph_result = tf_executor.graph {
    %output0, %control0 = tf_executor.island {
      %a = "tf.opA"(%arg0) : (tensor<i32>) -> tensor<i32>
      tf_executor.yield %a : tensor<i32>
    }
    %output1, %control1 = tf_executor.island {
      %b = "tf.opB"(%output0) : (tensor<i32>) -> tensor<i32>
      tf_executor.yield %b : tensor<i32>
    }
    tf_executor.fetch %output1 : tensor<i32>
  }

  std::unordered_set<string> outputs;
  TF_ASSERT_OK(FreezeSavedModel(saved_model_bundle, &frozen_graph_def, &inputs,
                                &outputs));
  std::unordered_set<string> expected_inputs = {"input0:0", "input1:0"};
  std::unordered_set<string> expected_outputs = {"output0:0", "output1:0"};
  EXPECT_EQ(expected_inputs, inputs);
  EXPECT_EQ(expected_outputs, outputs);
}

    value {
      s: "cluster"
    }
  }
}
node {
  name: "output0"
  op: "TPUReplicatedOutput"
  input: "Identity"
  attr {
    key: "T"
    value {
      type: DT_INT32
    }
  }
  attr {
    key: "num_replicas"
    value {
      i: 1
    }
  }
}
node {
  name: "output1"
  op: "TPUReplicatedOutput"
  input: "Identity_1"
  attr {
    key: "T"

// WholeModel-DAG: "tf.DumpTensor"(%[[output1]]) <{enabled = true, file_name = "unquantized_tensor_data.pb", func_name = "conv", log_dir_path = "/tmp/dumps/composite_conv2d_with_bias_and_relu6_fn_1", node_name = "Conv2D_1"}> : (tensor<*xf32>) -> ()
// WholeModel-DAG: return %[[output0]], %[[output1]]

// IntPerLayer-LABEL: func @conv
// IntPerLayer-DAG: %[[w:.*]] = "tf.Const"() <{value = dense<{{\[\[\[\[}}1.600000e-01, 1.000000e-01

  func.func @main() {
// CHECK: arith.constant
// CHECK: TPUCompile
// CHECK: TPUExecute
// CHECK-NOT: func @_func
    tf_executor.graph {
      %outputs, %control = tf_executor.island wraps "arith.constant"() {value = dense<2.000000e+00> : tensor<f32>} : () -> tensor<f32>
      %outputs_0, %control_1 = tf_executor.island wraps "arith.constant"() {value = dense<3.000000e+00> : tensor<f32>} : () -> tensor<f32>

    %outputs_0, %control_0 = tf_executor.island wraps "tf.Placeholder"() {device = "", dtype = "tfdtype$DT_FLOAT", name = "y", shape = "tfshape$dim { }"} : () -> tensor<0xf32>
    %outputs_1, %control_1 = tf_executor.island wraps "tf.TPUReplicatedInput"(%outputs_0) {N = 1 : i64, T = "tfdtype$DT_FLOAT", device = "", name = "input1"} : (tensor<0xf32>) -> tensor<0xf32>

}

func.func @not_modified(%arg0: tensor<f32>, %arg1: tensor<1x224x224x3xf32>) -> (tensor<1x401408xf32>, tensor<1x224x224x3xf32>) attributes {tf.entry_function = {control_outputs = "", inputs = "input0,input1", outputs = "output0,output1"}} {
  %cst = arith.constant dense<[1, 401408]> : tensor<2xi32>