// result.
  void Finalize();

  SmallVector<BlockArgument, 4> GetArgs() { return args_; }

  llvm::DenseMap<std::pair<mlir::Operation*, int>, int> GetResultStates() {
    return result_states_;
  }

  DenseMap<OpWithResultIndex, QuantStateIndex> result_states_;

  // Returns the state of the block argument.
  QuantState& GetArgQuantState(BlockArgument arg) {
    return states_[arg_states_[arg]];
  }

// "inputs" key.
BlockArgument CreateFilePrefixArg(func::FuncOp main_func_op) {
  Builder builder(main_func_op);

  // Add a new argument of type `tensor<!tf_type.string>` and update the
  // function type.
  auto file_prefix_arg_type =
      RankedTensorType::get(/*shape=*/{}, builder.getType<TF::StringType>());
  BlockArgument new_file_prefix_arg =

};

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

    unsigned GetNumPackedBlockArguments();
    llvm::ArrayRef<BlockArgument> GetPackedBlockArguments();
    llvm::ArrayRef<BlockArgument> GetReplicatedBlockArguments();
    bool IsReplicatedBlockArgument(BlockArgument block_arg);
    bool IsPackedBlockArgument(BlockArgument block_arg);
    unsigned GetReplicaOperandIndexForBlockArgument(BlockArgument block_arg, unsigned replica);

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

    if (tpose_arg1_type.getRank() != tpose_arg2_type.getRank()) {
      return failure();
    }

    if (llvm::isa<BlockArgument>(tpose_arg1.getPerm()) ||
        llvm::isa<BlockArgument>(tpose_arg2.getPerm())) {
      return failure();
    }
    auto perm1 = llvm::dyn_cast_or_null<arith::ConstantOp>(
        tpose_arg1.getPerm().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();
}