tf_device.return %2 : tensor<2xi32>
    }

    func.return %1 : tensor<2xi32>
  }

  // Tests extraction of an outside compiled tf.IfRegion op where the entirety
  // of tf.IfRegion op is outside compiled and wrapped inside another
  // tf.IfRegion op

  // CHECK-LABEL: func @outside_compiled_tf_if_nested
  func.func @outside_compiled_tf_if_nested(%arg0: tensor<2xi32>) -> tensor<2xi32> {

                                  {"tf.Yield"(%if12) : (!tf_variant) -> ()})
        {is_stateless = true} : (tensor<i1>) -> !tf_variant
    %if14 = "tf.IfRegion"(%arg0) ({"tf.Yield"(%if13) : (!tf_variant) -> ()},
                                  {"tf.Yield"(%if13) : (!tf_variant) -> ()})
        {is_stateless = true} : (tensor<i1>) -> !tf_variant

// -----

// Test invalid tf.Yield operation (parent should be IfRegion)
func.func @testInvalidYieldOp(%arg0: f32) -> () {
  // expected-error @+1 {{'tf.Yield' op expects parent op to be one of 'tf.CaseRegion, tf.IfRegion, tf.WhileRegion, tf.GeneratorDatasetRegion'}}
  "tf.Yield"(%arg0) : (f32) -> ()
}

// -----

// Test valid tf.IfRegion operation
// CHECK-LABEL: func @testValidIfRegionOp

  let summary = "Merges IfRegion ops together with a common predicate.";

  let description = [{
    This pass merges IfRegion ops together if they have the same predicate and it
    is safe to do so (there are no intermediate dependencies, they are in the
    same block, etc).

    For example:

    ```mlir
    "tf.IfRegion"(%0) ( {
      %2 = "tf.A"() : () -> (tensor<f32>)

  // However, its probably better to fold the IfRegion instead of having the
  // dead branch stay.

  // Inline the region in place of the IfRegion op, and forward the yield
  // inputs to the IfRegion op results. This is possible only if the yield
  // types match the result types.
  auto yield = cast<YieldOp>(region.front().getTerminator());
  auto updated_results = llvm::to_vector<4>(yield.getOperands());

  func.return %graph : tensor<i1>
  // expected-remark@above {{ID: 4}}
  // expected-remark@above {{Sinks: {}}}
}

// -----

// Tests that the pass does not add control dependencies for a stateless
// IfRegion op.

// CHECK-LABEL: func @stateless_ifregion_op
func.func @stateless_ifregion_op(
  // expected-remark@above {{ID: 18}}
  %arg0: tensor<*x!tf_type.resource<tensor<32xf32>>>,
  %arg1: tensor<i1>) {

  %true = "tf.Const"() {value = dense<true> : tensor<i1>} : () -> tensor<i1>

  // CHECK: [[Val0:%.*]] = "tf.Mul"(%arg0, %arg1)
  // CHECK-SAME: device = "/job:localhost/replica:0/task:0/device:GPU:0"
  %0 = "tf.IfRegion"(%true) ({
      %true_value = "tf.Mul"(%arg0, %arg1) {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<f32>, tensor<f32>) -> tensor<f32>
      "tf.Yield"(%true_value) : (tensor<f32>) -> ()
    }, {

  bool InferShapeForRestore(Operation* op);

  // Infers the shape IfOp outputs based on the shapes of the then and else
  // function result types.
  bool InferShapeForIf(IfOp op);

  // Infers the shape IfRegion outputs based on the shapes of the then and else
  // yields.
  bool InferShapeForIfRegion(IfRegionOp op);

  // Infers the shape CaseOp outputs based on the shapes of branch function
  // result types.

  func.func @cluster(%arg0 : tensor<i1>) {
    "tf_device.parallel_execute"() ({
      "tf.IfRegion"(%arg0) ({
        // expected-error@+1 {{The ParallelExecute ancestor of a ClusterFunc must be its direct parent.}}