// ReplicateOp) as the last operand of TPUReshardVariablesOp.
  BlockArgument state = unformat_replicate.GetReplicatedBlockArguments().back();
  auto replicated_block_args =
      unformat_replicate.GetReplicatedBlockArguments().drop_back(1);
  auto packed_block_args = unformat_replicate.GetPackedBlockArguments();
  unformat_operands.append(replicated_block_args.begin(),

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

}

// Returns the replicated block arguments. A copy should be made if the
// 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.