%elem_2 = "tf.Const"() {value = dense<10.0> : tensor<9xf32>} : () -> tensor<9xf32>
    %tl_set_item_2 = "tf.TensorListSetItem"(%tl, %idx, %elem_2) : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>, tensor<9xf32>) -> tensor<!tf_type.variant<tensor<*xf32>>>
    %elem_shape_3 = "tf.Const"() {value = dense<-1> : tensor<1xi32>} : () -> tensor<1xi32>

  # CHECK-SAME:   %arg0: tensor<1xf32> {tf._user_specified_name = "d", tf_saved_model.index_path = [0, "a"]},
  # CHECK-SAME:   %arg1: tensor<2xf32> {tf._user_specified_name = "d", tf_saved_model.index_path = [0, "b"]},
  # CHECK-SAME:   %arg2: tensor<3xf32> {tf._user_specified_name = "d", tf_saved_model.index_path = [0, "c"]},
  # CHECK-SAME:   %arg3: tensor<4xf32> {tf._user_specified_name = "d", tf_saved_model.index_path = [0, "x"]},

module {
  // CHECK-LABEL: main
  func.func @main(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> {
    %0 = "tfl.squared_difference"(%arg0, %arg1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
    func.return %0 : tensor<4xf32>
    // CHECK:  [[VAL_0:%.*]] = tfl.sub %arg0, %arg1 {fused_activation_function = "NONE"} : tensor<4xf32>

  %2 = "tf.Relu6"(%1) {device = ""} : (tensor<*xf32>) -> tensor<*xf32>
  func.return %2 : tensor<*xf32>
}

// CHECK-LABEL: func @conv_with_non_constant_filter
// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<2xf32>}> : () -> tensor<2xf32>
// CHECK: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"(%arg0, %arg1, %[[CONST_0]])
// CHECK-SAME: f = @composite_conv2d_with_bias_and_relu6_fn_1

    return
  }
  // CHECK-LABEL: func private @foo.impl
  func.func private @foo.impl(%arg0: tensor<1xf32>, %arg1: tensor<2xf32>) -> (tensor<2xf32>, tensor<1xf32>) {
    return %arg1, %arg0 : tensor<2xf32>, tensor<1xf32>
  }

      // Some loop invariant math
      %li0 = "tf.Div"(%a, %b) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
      %cst = arith.constant dense<22.0> : tensor<f32>
      %li1 = "tf.Mul"(%li0, %cst) : (tensor<4xf32>, tensor<f32>) -> tensor<4xf32>

      %final = "tf.Add"(%add, %li1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
      "tf.Yield"(%final, %sub) : (tensor<4xf32>, tensor<i32>) -> ()
    }

// CHECK-LABEL: import_stats_skip
func.func @import_stats_skip(%arg0: tensor<4xf32>, %cst: tensor<i32>) -> (tensor<2xf32>,tensor<2xf32>) {
  %0:2 = "tfl.split"(%cst, %arg0) {num_splits = 2 : i32} : (tensor<i32>, tensor<4xf32>) -> (tensor<2xf32>, tensor<2xf32>)
    loc(fused["skip1", "skip2.cc":10:8, callsite("op" at "skip3.cc":10:8)])
  func.return %0#0, %0#1 : tensor<2xf32>, tensor<2xf32>

// CHECK-NEXT: "tfl.split"
// CHECK-NEXT: return
}

}

// CHECK-LABEL: add
func.func @add(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<1xf32> {
  %0 = "tf.AddV2"(%arg0, %arg1) : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
  func.return %0: tensor<1xf32>
// CHECK: %[[ADD_0:.*]] = "tf.AddV2"(%arg0, %arg1) : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
// CHECK: return %[[ADD_0]] : tensor<1xf32>
}

// CHECK-LABEL: softmax

    %a:2, %control1 = tf_executor.island wraps "tf.A"() {_class = ["loc:@class"]} : () -> (tensor<2xf32>, tensor<2xf32>)
    %s:2, %control2 = tf_executor.island wraps "tf.Split"(%c, %a#1) {num_split = 2 : i32} : (tensor<i32>, tensor<2xf32>) -> (tensor<1xf32>, tensor<1xf32>)
    tf_executor.fetch
  }
  func.return
}

// -----

// CHECK-LABEL: func @no_colocate_split_has_device

// CHECK: tf.MyCustomOp(tensor<4xf32>, tensor<4xf32>) -> (tensor<4xf32>, tensor<3xf32>) : {name = "MyCustomOp"}

func.func @main(tensor<4xf32>) -> tensor<4xf32> {
^bb0(%arg0: tensor<4xf32>):
  %0 = "tfl.pseudo_const" () {name = "Const", value = dense<1.0> : tensor<4xf32>} : () -> tensor<4xf32>
  %1 = "tfl.mul"(%arg0, %0) {fused_activation_function = "NONE", name = "mul"} : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>