%2 = "tf.Cast"(%1) {Truncate = false} : (tensor<1x2xbf16>) -> tensor<1x2xf32>
  %3 = "tf.IdentityN"(%2) {device = ""} : (tensor<1x2xf32>) -> tensor<1x2xf32>
  return %3 : tensor<1x2xf32>
}

// CHECK: func @cast_bf16_matmul_to_fp32
// CHECK-DAG: %[[cst:.*]] = "tf.Const"() <{value = dense<{{.*}}> : tensor<10x2xf32>}> : () -> tensor<10x2xf32>
// CHECK: %[[matmul:.*]] = "tf.MatMul"(%arg0, %[[cst]])

  // CHECK: return %[[RESULT]]
  func.return %0: tensor<10x2xf32>
}

// -----

func.func @select_first(%arg0: tensor<10x2xf32>, %arg1: tensor<10x2xf32>) -> tensor<10x2xf32> {
  func.return %arg0: tensor<10x2xf32>
}

// CHECK-LABEL: testMultiInputLegacyCallOp
func.func @testMultiInputLegacyCallOp(%arg0: tensor<10x2xf32>, %arg1: tensor<10x2xf32>) -> tensor<10x2xf32> {

// CHECK-DAG:       [[VAL_21:%.*]] = arith.constant...

func.func @no_pushdown_multiple_outputs(%arg0: tensor<1000x2x!quant.uniform<i8:f32, 7.812500e-03>>) -> tensor<1000xf32> {
  %0 = "tfl.dequantize"(%arg0) : (tensor<1000x2x!quant.uniform<i8:f32, 7.812500e-03>>) -> tensor<1000x2xf32>
  %1:2 = "tfl.unpack"(%0) {axis = 1 : i32, num = 2 : i32} : (tensor<1000x2xf32>) -> (tensor<1000xf32>, tensor<1000xf32>)
  func.return %1#0 : tensor<1000xf32>

// CHECK-NEXT: tfl.dequantize
// CHECK-NEXT: tfl.unpack
}

func.func private @batched_function(%arg0: tensor<1x3xf32>, %arg1: tensor<*x!tf_type.resource>) -> tensor<1x3xf32> {
  %0 = "tf.Identity"(%arg0) : (tensor<1x3xf32>) -> tensor<1x3xf32>
  func.return %0 : tensor<1x3xf32>
}

// -----

func.func @test_batch_function_with_invalid_symbol(%arg0: tensor<1x3xf32>, %arg1: tensor<!tf_type.resource<tensor<1x3xf32>>>) -> () {

// CHECK-LABEL: fuseTileWithBinaryOp
func.func @fuseTileWithBinaryOp(%arg0: tensor<1x1xf32>) -> tensor<1x2xf32> {
  %cst = arith.constant dense<[[1,2]]> : tensor<1x2xi32>
  %cst1 = arith.constant dense<[[3.0, 4.0]]> : tensor<1x2xf32>
  %0 = "tfl.sqrt"(%arg0) : (tensor<1x1xf32>) -> tensor<1x1xf32>
  %1 = "tfl.tile"(%0, %cst) : (tensor<1x1xf32>, tensor<1x2xi32>) -> tensor<1x2xf32>

  func.func private @uniform_quantize_0(%arg0: tensor<1x2xf32>, %arg1: tensor<1x1xf32>, %arg2: tensor<1x1xi8>) -> tensor<1x2xi32> {
    %0 = stablehlo.convert %arg0 : (tensor<1x2xf32>) -> tensor<1x2xi32>
    return %0 : tensor<1x2xi32>
  }
// CHECK: @uniform_quantize_0
  func.func private @uniform_quantize_1(%arg0: tensor<1x3xf32>, %arg1: tensor<1x1xf32>, %arg2: tensor<1x1xi8>) -> tensor<1x3xi8> {

func.func @add_dense_dense_float_mixfng_1_n() -> tensor<2x2xf32> {
  %cst_0 = arith.constant dense<[[1.5, -2.5]]> : tensor<1x2xf32>
  %cst_1 = arith.constant dense<[[-3.], [4.]]> : tensor<2x1xf32>

  %0 = "tfl.add"(%cst_0, %cst_1) {fused_activation_function = "NONE"} : (tensor<1x2xf32>, tensor<2x1xf32>) -> tensor<2x2xf32>

  func.return %0 : tensor<2x2xf32>

func.func @no_change_to_dag(%arg0: tensor<!tf_type.resource<tensor<10x3xf32>>>, %arg1: tensor<!tf_type.resource<tensor<10x3xf32>>>, %arg2: tensor<!tf_type.resource<tensor<10x3xf32>>>, %arg3: tensor<!tf_type.resource<tensor<10x3xf32>>>) -> (tensor<!tf_type.resource<tensor<10x3xf32>>>, tensor<!tf_type.resource<tensor<10x3xf32>>>, tensor<!tf_type.resource...

  // CHECK: return [[RI]]
  func.return %ri : tensor<!tf_type.resource<tensor<10x3xf32>>>
}

// -----

// CHECK-LABEL:func @partitioned_2d
// CHECK-SAME: ([[ARG0:%.*]]: tensor<10x3xf32>, [[ARG1:%.*]]: tensor<10x3xf32>)
func.func @partitioned_2d(%arg0: tensor<10x3xf32>, %arg1: tensor<10x3xf32>) -> tensor<20x3xf32> {
  // CHECK: [[PI_0:%.*]] = "tf.TPUPartitionedInputV2"([[ARG0]], [[ARG1]])