func.func private @f(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
  "SomeOp"(%arg1) : (tensor<f32>) -> ()
  return %arg1 : tensor<f32>
}

// CHECK-LABEL: removes_first_arg
func.func @removes_first_arg(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
  // CHECK: call{{.*}}(%arg1)
  %1 = func.call @f(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
  return %1 : tensor<f32>
}

func.func @test_batch_matmul_to_einsum(%arg0: tensor<1x2x3xf32>, %arg1: tensor<3x4xf32>) -> tensor<1x2x4xf32> {
  // CHECK-LABEL: test_batch_matmul_to_einsum
  // CHECK: "tf.Einsum"(%arg0, %arg1) <{equation = "...mk,...kn->...mn"}> : (tensor<1x2x3xf32>, tensor<3x4xf32>) -> tensor<1x2x4xf32>
  %0 = "tf.BatchMatMul"(%arg0, %arg1) {adj_x = false, adj_y = false} : (tensor<1x2x3xf32>, tensor<3x4xf32>) -> tensor<1x2x4xf32>

  %arg0: tensor<1x28x28x64xf32>,
  %arg1: tensor<64xf32>
) -> tensor<1x28x28x64xf32> {

  // CHECK: "tf.FusedBatchNormV3"
  // CHECK-SAME: (%[[X_TRANSPOSE:[0-9]*]], %arg1, %arg1, %arg1, %arg1)
  // CHECK-SAME: data_format = "NCHW"
  %y, %batch_mean, %batch_var, %reserve_1, %reserve_2, %reserve_3
    = "tf.FusedBatchNormV3"(%arg0, %arg1, %arg1, %arg1, %arg1)
       {
         data_format = "NHWC",

// CHECK-LABEL: listReserveScalarShapeI32
func.func @listReserveScalarShapeI32(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<!tf_type.variant<tensor<*xi32>>> {
  %0 = "tf.TensorListReserve"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<*xi32>>>

  // CHECK: tac.cost = 1.000000e+03
  %0 = "tfl.add"(%arg0, %arg1) {fused_activation_function = "RELU", tac.device = "CPU"} : (tensor<10x10x10xf32>, tensor<10xf32>) -> tensor<10x10x10xf32>
  // CHECK: tac.cost = 1.000000e+03

}

func.func @testAdd(%arg0: tensor<10x10x10xf32>, %arg1: tensor<10x10x10xf32>) -> tensor<10x10x10xf32> {
  // CHECK: _arithmetic_count = 1000 : i64
  %0 = "tfl.add"(%arg0, %arg1) {fused_activation_function = "NONE"} : (tensor<10x10x10xf32>, tensor<10x10x10xf32>) -> tensor<10x10x10xf32>
  func.return %0 : tensor<10x10x10xf32>
}

func.func @testAddBroadcast(%arg0: tensor<10x10x10xf32>, %arg1: tensor<10xf32>) -> tensor<10x10x10xf32> {

  // CHECK-NEXT: "tf.AssignVariableOp"([[ARG1]], [[PARALLEL_EXECUTE]]#0)
  "tf.AssignVariableOp"(%arg1, %0#0) : (tensor<!tf_type.resource>, tensor<i1>) -> ()
  // CHECK-NEXT: return [[PARALLEL_EXECUTE]]#0
  func.return %0#0 : tensor<i1>
}

// CHECK-LABEL: func @not_assign_var
// CHECK-SAME:  ({{.+}}: tensor<i1>, [[ARG1:%.+]]: tensor<!tf_type.resource>)
func.func @not_assign_var(%arg0: tensor<i1>, %arg1: tensor<!tf_type.resource>) {

  // CHECK: %[[MUL:.*]] = "tf.Mul"
  %0 = "tf.Mul"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
  // CHECK: %[[RESULT:.*]]:4 = "tf_device.cluster"
  %1:4 = "tf_device.cluster"() ({
    // CHECK: %[[ADD:.*]] = "tf.AddV2"
    %2 = "tf.AddV2"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>

CHECK:  %[[LSTM:[a-z0-9]*]] = "tfl.unidirectional_sequence_lstm"(%arg2, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg1, %arg1, %arg1, %arg1, %arg1, %arg1, %arg1, %arg0, %arg1, %[[CST_0]], %[[CST_1]], %arg0, %arg0, %arg0, %arg0) <{fused_activation_function = "NONE", time_major = true}> : (tensor<4x4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<4xf32>,...

  %arg0: tensor<1x28x28x64xf32>,
  %arg1: tensor<64xf32>
) -> tensor<1x28x28x64xf32> {

  // CHECK: "tf.FusedBatchNormV3"
  // CHECK-SAME: (%[[X_TRANSPOSE:[0-9]*]], %arg1, %arg1, %arg1, %arg1)
  // CHECK-SAME: data_format = "NCHW"
  %y, %batch_mean, %batch_var, %reserve_1, %reserve_2, %reserve_3
    = "tf.FusedBatchNormV3"(%arg0, %arg1, %arg1, %arg1, %arg1)
       {
         data_format = "NHWC",