%0 = stablehlo.constant dense<[-0.211145893, -0.708605706]> : tensor<2xf32>
        %1 = stablehlo.constant dense<[[-0.630731344, 0.54962182], [0.180364341, -0.764542698]]> : tensor<2x2xf32>

    }
  }
  if_op.replaceAllUsesWith(new_if);
  if_op.erase();
  return success();
}

// Handles stack usage by a tf.StatefulPartitionedCall or a tf.PartitionedCall.
// It will first check if the callee was previously handled, and try to reuse
// that result if so. Otherwise, it will clone and convert the callee function,
// and performs stack ops decomposition on it.

  %1 = "tf._SomeOp"() : () -> tensor<2xf32>
  "tf.AssignVariableOp"(%0, %1) : (tensor<!tf_type.resource<tensor<2xf32>>>, tensor<2xf32>) -> ()
  %2 = "tf.PartitionedCall"(%0) {config = "", config_proto = "", executor_type = "", f = @callee} : (tensor<!tf_type.resource<tensor<2xf32>>>) -> tensor<2xf32>
  func.return %2 : tensor<2xf32>
}

            MHLO_Tensor:$on_false),
           (CHLO_BroadcastSelectOp $pred, $on_true, $on_false)>;

//===----------------------------------------------------------------------===//
// PartitionedCall and LegacyCall op patterns.
//===----------------------------------------------------------------------===//

def ArgTypesMatchCallee : Constraint<

  %0 = "tf.Const"() {_output_shapes = ["tfshape$"], device = "", dtype = f32, value = dense<0.000000e+00> : tensor<f32>} : () -> tensor<f32>

  // CHECK-NEXT: [[results_1:%.*]] = call @callee(
  // CHECK-SAME: (!tf_mlrt.tensor) -> !tf_mlrt.tensor
  %2 = "tf.PartitionedCall"(%0) {config = "", config_proto = "", executor_type = "", f = @callee} : (tensor<i32>) -> (tensor<i32>)
  // CHECK-NEXT: [[results_2:%.*]] = call @callee(
  // CHECK-SAME: (!tf_mlrt.tensor) -> !tf_mlrt.tensor

    ```

    `_XlaMustCompile=true` in the following code

    ```mlir

```

`_XlaMustCompile=true` in the following code

```mlir

//===--------------------------------------------------------------------===//
//  tf.{|Stateful}PartitionedCall
//===--------------------------------------------------------------------===//

// Test valid tf.PartitionedCall
// CHECK-LABEL: func @testValidPartitionedCall
func.func @testValidPartitionedCall(%arg0: tensor<i32>) -> tensor<i32> {

  NameAttrList b_name_attr;
  b_name_attr.set_name("CompilableFn");
  ops::PartitionedCall b(root.WithOpName("B"), {a, a}, {DT_FLOAT}, b_name_attr);
  NameAttrList c_name_attr;
  c_name_attr.set_name("UncompilableFn");

  ops::PartitionedCall c(root.WithOpName("C"), {a}, {DT_FLOAT}, c_name_attr);
  Output d = ops::Add(root.WithOpName("D"), b.output.front(), c.output.front());