};

void MoveResourceArgsToEnd(func::FuncOp callee) {
  llvm::DenseMap<BlockArgument, BlockArgument> mapping;
  unsigned num_params = callee.getNumArguments();
  llvm::BitVector removed_params(num_params);
  // Copy the resource-type parameters to the end.
  for (unsigned i = 0; i < num_params; ++i) {
    BlockArgument param = callee.getArgument(i);
    if (mlir::isa<TF::ResourceType>(getElementTypeOrSelf(param.getType()))) {

static LogicalResult traceUpwardsToArgument(Value v, llvm::DenseSet<Value> seen,
                                            BlockArgument *out) {
  if (seen.contains(v)) {
    return failure();  // infinite loop
  }
  seen.insert(v);

  if (auto blockArg = mlir::dyn_cast<BlockArgument>(v)) {
    Operation *op = blockArg.getOwner()->getParentOp();

    // If we're in the first block, then the argument to that block is the

      });

  py::class_<mlir::Value>(m, "Value").def("getType", &mlir::Value::getType);
  py::class_<mlir::OpResult, mlir::Value>(m, "OpResult");
  py::class_<mlir::BlockArgument, mlir::Value>(m, "BlockArgument");

ResourceConstructingOps ResourceConstructingOps::EntryState(
    MLIRContext *context) {
  return ResourceConstructingOps();
}
ResourceConstructingOps ResourceConstructingOps::EntryState(Value value) {
  if (auto barg = mlir::dyn_cast<BlockArgument>(value)) {
    if (func::FuncOp func =
            dyn_cast<func::FuncOp>(barg.getOwner()->getParentOp())) {
      SymbolTable symbol_table(func->getParentOfType<ModuleOp>());

  return op->getAttrOfType<StringAttr>(kXlaOutsideCompilationAttr) != nullptr;
}

// Finds op that created a given value. If the value is a BlockArgument, this
// returns the owner of the Block.
Operation* GetOpOfValue(Value value) {
  if (auto block_arg = mlir::dyn_cast<BlockArgument>(value))
    return block_arg.getOwner()->getParentOp();

  return value.getDefiningOp();
}

          TPUAnnotateDynamicShapeInputsPass> {
  void runOnOperation() override;
};

// Finds op that created a given value. If the value is a BlockArgument, this
// returns the owner of the Block.
Operation* GetOpOfValue(Value value) {
  if (auto block_arg = mlir::dyn_cast<BlockArgument>(value))
    return block_arg.getOwner()->getParentOp();

  return value.getDefiningOp();
}

          while_op, while_op.getBodyAttr());
      for (auto &arg : while_op->getOpOperands()) {
        BlockArgument cond_arg = cond.getArgument(arg.getOperandNumber());
        this->join(getLatticeElement(cond_arg), *getLatticeElement(arg.get()));
        BlockArgument body_arg = body.getArgument(arg.getOperandNumber());
        this->join(getLatticeElement(body_arg), *getLatticeElement(arg.get()));
      }

    }
  }

  // Collect all those freezable. This is an extra scan but allows for the
  // partial behavior from `allow_mutable_tensor`.
  DenseMap<BlockArgument, bool> freezeable;
  for (auto func : module.getOps<func::FuncOp>()) {
    for (BlockArgument val : func.getArguments()) {
      if (!mlir::isa<TF::ResourceType>(getElementTypeOrSelf(val.getType())))
        continue;

    // defining op.
    if (def_op != nullptr &&
        (isa<TF::TPUExecuteOp, TF::TPUExecuteAndUpdateVariablesOp>(def_op))) {
      return true;
    }
    if (BlockArgument block_arg = mlir::dyn_cast<BlockArgument>(arg)) {
      // Skip the folding logic if the block argument is not from the function
      // arguments. This can happen when the argument is from a while loop.

                           result.getResultNumber());
    return std::string(result.getOwner()->getName().getStringRef());
  }
  // Use the ASM syntax for BlockArgument
  if (auto arg = mlir::dyn_cast<mlir::BlockArgument>(val)) {
    return "arg" + std::to_string(arg.getArgNumber());
  }
  return "";
}

std::string OpOrArgStripNameMapper::GetName(OpOrVal op_or_val) {