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

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

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

  // Update input type.
  auto transform_result_type =
      RankedTensorType::get(new_shape, getElementTypeOrSelf(cast_input));
  cast_input.setType(transform_result_type);
  auto block_arg = mlir::dyn_cast<mlir::BlockArgument>(cast_input);
  auto cast_op_input = dyn_cast_or_null<TF::CastOp>(cast_input.getDefiningOp());
  while (block_arg || cast_op_input) {
    if (block_arg) {
      // Change on device function type/shape.

// `tf.ReadVariableOp` and `tf.AssignVariableOp` only.
mlir::LogicalResult ValidateResourceArgument(func::FuncOp function,
                                             BlockArgument resource_arg,
                                             TF::ResourceType resource_type) {
  if (resource_type.getSubtypes().size() != 1)
    return function.emitError()
           << "expects resource type of argument "

    return false;
  llvm::SmallSetVector<Value, 8> aliases =
      resource_alias_analysis.GetResourceAliases(resource_iterator);

  auto is_generator = [](Value val) {
    if (mlir::isa<BlockArgument>(val)) return true;
    Operation* definition = val.getDefiningOp();
    return definition->getNumOperands() == 0 &&
           definition->getNumResults() == 1;
  };

// replicate op is being modified.
llvm::ArrayRef<BlockArgument> ReplicateOp::GetReplicatedBlockArguments() {
  return GetBody().getArguments().drop_back(GetNumPackedBlockArguments());
}

// Returns the packed block arguments. A copy should be made if the replicate op
// is being modified.
llvm::ArrayRef<BlockArgument> ReplicateOp::GetPackedBlockArguments() {

                      dst_index);
    }
    return absl::OkStatus();
  }

  auto input_arg = mlir::cast<BlockArgument>(src);
  auto input_node_it = args_.find(input_arg);
  TF_RET_CHECK(input_node_it != args_.end())
      << "Use of BlockArgument encounted before def!";
  // For argument, there is only one result output, so the index is always 0.

                      dst_index);
    }
    return absl::OkStatus();
  }

  auto input_arg = mlir::cast<BlockArgument>(src);
  auto input_node_it = args_.find(input_arg);
  TF_RET_CHECK(input_node_it != args_.end())
      << "Use of BlockArgument encounted before def!";
  // For argument, there is only one result output, so the index is always 0.