#include "tensorflow/c/experimental/filesystem/modular_filesystem.h"
#include "tensorflow/cc/framework/gradients.h"
#include "tensorflow/cc/framework/ops.h"
#include "tensorflow/cc/framework/scope_internal.h"
#include "tensorflow/cc/ops/while_loop.h"
#include "tensorflow/cc/saved_model/loader.h"
#include "tensorflow/core/distributed_runtime/server_lib.h"
#include "tensorflow/core/framework/logging.h"
#include "tensorflow/core/framework/op_gen_lib.h"

  LogicalResult matchAndRewrite(WhileRegionOp while_op,
                                PatternRewriter &rewriter) const override {
    auto &body_block = while_op.getBody().front();
    auto &cond_block = while_op.getCond().front();
    bool changed = false;
    for (int op_idx : llvm::seq<int>(0, while_op.getNumOperands())) {
      auto body_arg = body_block.getArgument(op_idx);

          while_op.getInput().getTypes(), while_op.getOutput().getTypes(),
          while_op.getBody().getArgumentTypes());
      return PropagateShapeToRegions(compatible_types,
                                     {&while_op.getCond(), &while_op.getBody()},
                                     max_iterations);
    }
    return PropagateShapeToRegions(while_op.getInput().getTypes(),

    SmallVector<Value, 4> new_body_yield;
    SmallVector<bool, 4> removed_operand(while_op.getNumOperands(), false);
    llvm::SmallVector<Type, 4> types;
    new_operands.reserve(while_op.getNumOperands());
    new_body_yield.reserve(while_op.getNumOperands());
    types.reserve(while_op.getNumOperands());

    // Remove block arguments not used in either cond or body. This leaves the

    auto new_while = rewriter.create<TF::WhileRegionOp>(
        while_op.getLoc(), while_op->getResultTypes(), while_op->getOperands(),
        /*parallel_iterations=*/10,
        /*is_stateless=*/false, /*shape_invariant=*/false);
    new_while.getCond().takeBody(while_op.getCond());
    new_while.getBody().takeBody(while_op.getBody());
    ReplaceReturnOp(new_while.getCond(), rewriter);