decomposed_partitioned_call_callees))) {
        return failure();
      }
    } else if (llvm::isa<TF::CaseRegionOp>(op) ||
               llvm::isa<TF::IfRegionOp>(op) ||
               llvm::isa<TF::WhileRegionOp>(op)) {
      if (failed(HandleRegionControlFlowOps(
              op, module, stats, decomposed_partitioned_call_callees)))
        return failure();

  // 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.
  bool InferShapeForCase(CaseOp op);

  // Infers the shape CaseRegion outputs based on the shapes of the branch

                                         vars_initialized,
                                         lifted_partitioned_call_callees))) {
        return failure();
      }
    } else if (isa<TF::IfRegionOp, TF::CaseRegionOp, TF::WhileRegionOp>(op)) {
      for (Region& region : op.getRegions())
        (void)HoistForControlFlow(&region.front(), module, vars_initialized,

    }
  }];
}

def TF_YieldOp : TF_Op<"Yield",
      [Terminator,
       Pure,
       NativeOpTrait<"ReturnLike", [], "", "">,
       ParentOneOf<["CaseRegionOp", "IfRegionOp", "WhileRegionOp", "GeneratorDatasetRegionOp"]>,
       DeclareOpInterfaceMethods<RegionBranchTerminatorOpInterface,
           ["getMutableSuccessorOperands"]>,
      ]> {
  let summary = "Yield operation";

  LogicalResult matchAndRewrite(
      mhlo::IfOp op, OpAdaptor adaptor,
      ConversionPatternRewriter& rewriter) const final {
    // HLO IfOp currently doesn't support stateless
    auto new_op = rewriter.create<TF::IfRegionOp>(
        op.getLoc(), op->getResultTypes(), op.getPred(),
        /*is_stateless=*/false, /*_then_func_name=*/nullptr,
        /*_else_func_name=*/nullptr);
    new_op.getThenBranch().takeBody(op.getTrueBranch());

    ConvertSigmoidGradOpDynamic,
    ConvertConv2DDynamic,
    ConvertPadOpDynamic,
    ConvertGatherNdOpDynamic,
    LowerControlFlowOp<TF::CaseRegionOp, mhlo::CaseOp>,
    LowerControlFlowOp<TF::IfRegionOp, mhlo::IfOp>,
    LowerControlFlowOp<TF::WhileRegionOp, mhlo::WhileOp>,
    LowerYieldOp>(context);
  // clang-format on
}
// LINT.ThenChange(:MlirAlwaysOps)
}  // end namespace mhlo