}

constexpr absl::string_view kFunctionWithRegion =
    R"mlir(
  func.func @main(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
    %if = "stablehlo.if"(%arg0) ({
      %0 = stablehlo.add %arg1, %arg1 : tensor<f32>
      stablehlo.return %0 : tensor<f32>
    }, {
      %1 = stablehlo.add %arg2, %arg2 : tensor<f32>
      stablehlo.return %1 : tensor<f32>
    }) : (tensor<i1>) -> (tensor<f32>)

  func.return %id0, %id0, %arg1 : !tf_res, !tf_res, !tf_res
}

// expected-remark@below {{Region #0, Arg #0, ID 0 : 0}}
// expected-remark@below {{Region #0, Arg #1, ID 1 : 1}}
func.func @case_branch2(%arg0: !tf_res, %arg1: !tf_res) -> (!tf_res, !tf_res, !tf_res) {
  func.return %arg0, %arg0, %arg1 : !tf_res, !tf_res, !tf_res
}

// -----
// Test aliasing through WhileOp

// CHECK-SAME: (%[[ARG_0:[a-z0-9]*]]: tensor<*xf32>, %[[ARG_1:[a-z0-9]*]]: tensor<*xf32>, %[[ARG_2:[a-z0-9]*]]: tensor<*xf32>)
func.func @replicate_arg_shape_with_packed(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>, %arg2: tensor<*xf32>) {
  %0:4 = tf_device.replicate([%arg0, %arg1] as %ri: tensor<*xf32>, %arg2 as %rj : tensor<*xf32>) {n = 2: i32} {
    %1 = "tf.Shape"(%rj) {device = "", T = "tfdtype$DT_FLOAT", out_type = "tfdtype$DT_INT32"} : (tensor<*xf32>) -> tensor<?xi32>

// CHECK-SAME: f = @composite_conv2d_with_bias_and_relu6_fn_1}>
// CHECK-SAME: {_tfl_quant_trait = "fully_quantizable"
// CHECK: %[[PARTITIONEDCALL_1:.*]] = "tf.PartitionedCall"(%arg0, %arg1, %[[CONST_0]])
// CHECK-SAME: f = @composite_conv2d_with_bias_and_relu_fn_1}
// CHECK: %[[PARTITIONEDCALL_2:.*]] = "tf.PartitionedCall"(%arg0, %arg1, %[[CONST_0]])

def LegalizePadV2 : Pat<(TF_PadV2Op $arg0, $arg1, $cst),
                        (TFL_PadV2Op $arg0, $arg1, $cst)>;

def LegalizeMean : Pat<(TF_MeanOp $arg0, $arg1, BoolAttr:$arg2),
                       (TFL_MeanOp $arg0, $arg1, $arg2)>;

def LegalizeSum : Pat<(TF_SumOp $arg, $axes, BoolAttr:$arg2),
                      (TFL_SumOp $arg, (CreateTFCastToInt32Op $axes), $arg2)>;

// CHECK-NEXT: return %[[ES]] : tensor<3x2x3xf32>
}

// CHECK-LABEL: my_pack_multiple_inputs
func.func @my_pack_multiple_inputs(%arg0: tensor<2x3xf32>, %arg1: tensor<2x3xf32>, %arg2: tensor<2x3xf32>) -> tensor<3x2x3xf32> {
  %0 = "tf.MyPack"(%arg0, %arg1, %arg2) {N=3:i32, axis=0:i32} : (tensor<2x3xf32>, tensor<2x3xf32>, tensor<2x3xf32>) -> tensor<3x2x3xf32>
  func.return %0 : tensor<3x2x3xf32>

    %0 = "tf.MatMul"(%arg0, %arg1) <{grad_a = false, grad_b = false, transpose_a = false, transpose_b = false}> {attr_map = "0:transpose_a,1:transpose_b", device = ""} : (tensor<1x4xf32>, tensor<4x3xf32>) -> tensor<1x3xf32>

A#loc9 = loc(fused[#loc1, #loc2, #loc3])\0Amodule {\0A  func.func @host_func(%arg0: tensor<3360x?xi32> loc(fused[#loc1, #loc2, #loc3]), %arg1: tensor<3xi32> loc(fused[#loc1, #loc2, #loc3]), %arg2: tensor<i32> loc(fused[#loc1, #loc2, #loc3]), %arg3: tensor<2xi32> loc(fused[#loc1, #loc2, #loc3])) -> (tensor<1x1120x?xi32>, tensor<1x1120x?xi32>, tensor<1120x?xi32>, tensor<2xi32>) {\0A    %0 = \22tf.Reshape\22(%arg0, %arg1) {_xla_outside_compilation = \220\22} : (tensor<3360x?xi32>, tensor<3xi32>) ->...

}

// CHECK-LABEL: RemoveTrival
// QDQ-LABEL: RemoveTrival
func.func @RemoveTrival(%arg0: tensor<384x512x!quant.uniform<i8:f32, 1.0:-128>>, %arg1: tensor<128x512x!quant.uniform<i8<-127:127>:f32, 1.0>>, %arg2: none) -> tensor<384x128x!quant.uniform<i8:f32, 2.0>> {

}

// CHECK-LABEL: func @switch_with_control_inputs(
func.func @switch_with_control_inputs(%arg0: tensor<i1>, %arg1: !tf_executor.control, %arg2: !tf_executor.control) -> tensor<i1> {
  %result = tf_executor.graph {
// CHECK: tf_executor.Switch %{{[^%]*}}, %{{[^%]*}}, %{{[^%]*}}, %{{[^%]*}} : tensor<i1>
    %1:3 = tf_executor.Switch %arg0, %arg0, %arg1, %arg2 : tensor<i1>
    tf_executor.fetch %1#0 : tensor<i1>
  }
  func.return %result : tensor<i1>