// 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>) -> ()
    }

  func.return %add : tensor<*xf32>
}
// CHECK: func @decompose_batch_norm
// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() <{value = dense<2.49743462E-5> : tensor<2xf32>}> : () -> tensor<2xf32>

}

// -----

// CHECK-LABEL: @acos
func.func @acos(%arg0: tensor<2xf32>) -> tensor<2xf32> {
  // CHECK-NOT:  tf.Acos
  %0 = "tf.Acos"(%arg0) : (tensor<2xf32>) -> tensor<2xf32>
  func.return %0 : tensor<2xf32>
}

// -----

// NOFALLBACK-LABEL: @xla_svd
func.func @xla_svd(%arg0: tensor<1x1xf32>) -> (tensor<1xf32>, tensor<1x1xf32>, tensor<1x1xf32>) {
  // NOFALLBACK: XlaSvd

// 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>

    %1 = "tfl.abs"(%arg2) {fused_activation_function = "NONE"} : (tensor<5xf32>) -> tensor<5xf32>
    %2 = stablehlo.add %1, %arg2 : tensor<5xf32>
    %3 = "tfl.abs"(%2) {fused_activation_function = "NONE"} : (tensor<5xf32>) -> tensor<5xf32>
    "stablehlo.return"(%3) : (tensor<5xf32>) -> ()
  }) {dimensions = array<i64: 0>} : (tensor<7x5xf32>, tensor<5xf32>) -> tensor<5xf32>
  func.return %0: tensor<5xf32>
}

// -----

// CHECK-NEXT:  %[[quant:.*]] = "tfl.quantize"(%[[split]]#0) <{qtype = tensor<2x!quant.uniform<u8:f32, 1.000000e+00>>}> {volatile} : (tensor<2xf32>) -> tensor<2x!quant.uniform<u8:f32, 1.000000e+00>>
// CHECK-NEXT:  return %[[quant]] : tensor<2x!quant.uniform<u8:f32, 1.000000e+00>>

    %cst_0 = "tf.Const"() {value = dense<1.000000e+00> : tensor<2xf32>} : () -> tensor<2xf32>
    %var_0 = "tf.VarHandleOp"() {shared_name = "var_0"} : () -> tensor<!tf_type.resource<tensor<2xf32>>>
    "tf.AssignVariableOp"(%var_0, %cst_0) : (tensor<!tf_type.resource<tensor<2xf32>>>, tensor<2xf32>) -> ()
    %cst_1 = "tf.Const"() {value = dense<2.000000e+00> : tensor<3xf32>} : () -> tensor<3xf32>