func.func @multiple_return_no_controls_needed(%arg0: tensor<*xi32>, %arg1: tensor<i32>) -> (tensor<*xi32>, tensor<*xi32>) {
  %graph:2 = tf_executor.graph {
    %add1, %add1_control = tf_executor.island wraps "tf.Add"(%arg0, %arg1) : (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
    %add2, %add2_control = tf_executor.island wraps "tf.Add"(%add1, %arg1) : (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
    tf_executor.fetch %add1, %add2 : tensor<*xi32>, tensor<*xi32>

        %cst_0 = "tf.Const"() {device = "", value = dense<[3, 4, 5, 6]> : tensor<4xi32>} : () -> tensor<4xi32>
        %1 = "tf.VarHandleOp"() {container = "", device = "/device:CPU:0", shared_name = "var_1"} : () -> tensor<!tf_type.resource<tensor<4xi32>>>
        "tf.AssignVariableOp"(%1, %cst_0) : (tensor<!tf_type.resource<tensor<4xi32>>>, tensor<4xi32>) -> ()

        return
      }
    }
  )mlir";

    %cst_1 = "tf.Const"() {value = dense<2> : tensor<4xi32>} : () -> tensor<4xi32>
    %var_1 = "tf.VarHandleOp"() {shared_name = "var_1"} : () -> tensor<!tf_type.resource<tensor<4xi32>>>
    "tf.AssignVariableOp"(%var_1, %cst_1) : (tensor<!tf_type.resource<tensor<4xi32>>>, tensor<4xi32>) -> ()
    return
  }

// CHECK: func.func @init_func_restore_op_multiple_variables

// CHECK: %[[reshape_i8:.*]] = "tf.Reshape"(%[[maxpool_i8]]
// CHECK-SAME: (tensor<*xf32>, tensor<2xi32>) -> tensor<*xf32>

// -----

// CHECK-LABEL: standalone_avgpool_test
func.func @standalone_avgpool_test(%arg0: tensor<*xf32>) -> tensor<*xf32> {
  %cst = arith.constant dense<[-1, 144]> : tensor<2xi32>
  %0 = "quantfork.qcast"(%arg0) : (tensor<*xf32>) -> tensor<*x!quant.uniform<i8:f32, 5.000000e-02:-10>>

    %cst_1 = "tf.Const"() {value = dense<[2, 1024]> : tensor<2xi32>} : () -> tensor<2xi32>
    // Original identity preserved.
    %cst_identity = "tf.Identity"(%cst_0) {device = ""} : (tensor<1024x2xf32>) -> tensor<1024x2xf32>
    %0 = "tf.Reshape"(%cst_identity, %cst_1) : (tensor<1024x2xf32>, tensor<2xi32>) -> tensor<2x1024xf32>

            step_marker_location = ""} : (tensor<*xi32>) -> tensor<*xi32>
        return
      }
      func @func(%arg0: tensor<*xi32> {mhlo.sharding = "\01\02\03"}) ->
                (tensor<*xi32> {mhlo.sharding = "\08\01\1A\01\01\22\01\00"}) {
        %0 = "tf.XlaSharding"(%arg0) {_XlaSharding = "\01\02\03", sharding = "\01\02\03"} : (tensor<*xi32>) -> tensor<*xi32>
        %1 = "tf.A"(%0) : (tensor<*xi32>) -> tensor<*xi32>

  %zps_0 = "tf.Const"() { value = dense<[3, 4]> : tensor<2xi32> } : () -> tensor<2xi32>
  %scales_1 = "tf.Const"() { value = dense<[3.0, 4.0]> : tensor<2xf32> } : () -> tensor<2xf32>
  %zps_1 = "tf.Const"() { value = dense<[5, 6]> : tensor<2xi32> } : () -> tensor<2xi32>

  func.return %0 : tensor<1xi32>
}

// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1xi32> {tf_saved_model.index_path = ["input"]}) -> (tensor<1xi32> {tf_saved_model.index_path = ["r"]})
  attributes {tf_saved_model.exported_names = ["main"]} {

func.func @init() attributes {tf_saved_model.exported_names = ["__tf_saved_model_session_initializer_init"]} {
  %keys = "tf.Const"() {value = dense<[1, 2, 3, 4]> : tensor<4xi32> } : () -> tensor<4xi32>
  %values = "tf.Const"() {value = dense<[1, 2, 3, 4]> : tensor<4xi32> } : () -> tensor<4xi32>
  %handle = "tf.HashTableV2"() {container = "", key_dtype = i32, shared_name = "hash_table", use_node_name_sharing = false, value_dtype = i32} : () -> tensor<!tf_type.resource>

    %4 = "tf.LayerNorm"(%a1, %a2, %a3, %a4) {_tfl_quant_trait = "fully_quantizable", device = ""} : (tensor<128x128xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xi32>) -> tensor<128x128xf32>
   "tfl.yield"(%4) : (tensor<128x128xf32>) -> ()
  }) {_tfl_quant_trait = "fully_quantizable", device = ""} : (tensor<128x128xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xi32>) -> tensor<128x128xf32>