class FoldConstantIfRegionOp : public OpRewritePattern<TF::IfRegionOp> {
 public:
  explicit FoldConstantIfRegionOp(MLIRContext* context)
      : OpRewritePattern<TF::IfRegionOp>(context) {}
  LogicalResult matchAndRewrite(TF::IfRegionOp op,
                                PatternRewriter& rewriter) const override;
};

LogicalResult FoldConstantIfRegionOp::matchAndRewrite(
    TF::IfRegionOp op, PatternRewriter& rewriter) const {

  // 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

  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());