class_docstr: Datapipe designed to operate over binary (npy) versions of Criteo datasets. Loads
  the entire dataset into memory to prevent disk speed from affecting throughout. Each
  rank reads only the data for the portion of the dataset it is responsible for.
  
  The torchrec/datasets/scripts/npy_preproc_criteo.py script can be used to convert
  the Criteo tsv files to the npy files expected by this dataset.
  
  Args:
      stage (str): "train", "val", or "test".
      dense_paths (List[str]): List of path strings to dense npy files.
      sparse_paths (List[str]): List of path strings to sparse npy files.
      labels_paths (List[str]): List of path strings to labels npy files.
      batch_size (int): batch size.
      rank (int): rank.
      world_size (int): world size.
      shuffle_batches (bool): Whether to shuffle batches
      hashes (Optional[int]): List of max categorical feature value for each feature.
          Length of this list should be CAT_FEATURE_COUNT.
      path_manager_key (str): Path manager key used to load from different
          filesystems.
  
  Example::
  
      template = "/home/datasets/criteo/1tb_binary/day_{}_{}.npy"
      datapipe = InMemoryBinaryCriteoIterDataPipe(
          dense_paths=[template.format(0, "dense"), template.format(1, "dense")],
          sparse_paths=[template.format(0, "sparse"), template.format(1, "sparse")],
          labels_paths=[template.format(0, "labels"), template.format(1, "labels")],
          batch_size=1024,
          rank=torch.distributed.get_rank(),
          world_size=torch.distributed.get_world_size(),
      )
      batch = next(iter(datapipe))
class_docstr_tok: ["Datapipe","designed","to","operate","over","binary","(","npy",")","versions","of","Criteo","datasets",".","Loads","the","entire","dataset","into","memory","to","prevent","disk","speed","from","affecting","throughout",".","Each","rank","reads","only","the","data","for","the","portion","of","the","dataset","it","is","responsible","for",".","The","torchrec\/datasets\/scripts\/npy_preproc_criteo.py","script","can","be","used","to","convert","the","Criteo","tsv","files","to","the","npy","files","expected","by","this","dataset",".","Args",":","stage","(","str",")",":","``","train","''",",","``","val","''",",","or","``","test","''",".","dense_paths","(","List","[","str","]",")",":","List","of","path","strings","to","dense","npy","files",".","sparse_paths","(","List","[","str","]",")",":","List","of","path","strings","to","sparse","npy","files",".","labels_paths","(","List","[","str","]",")",":","List","of","path","strings","to","labels","npy","files",".","batch_size","(","int",")",":","batch","size",".","rank","(","int",")",":","rank",".","world_size","(","int",")",":","world","size",".","shuffle_batches","(","bool",")",":","Whether","to","shuffle","batches","hashes","(","Optional","[","int","]",")",":","List","of","max","categorical","feature","value","for","each","feature",".","Length","of","this","list","should","be","CAT_FEATURE_COUNT",".","path_manager_key","(","str",")",":","Path","manager","key","used","to","load","from","different","filesystems",".","Example",":",":","template","=","``","\/home\/datasets\/criteo\/1tb_binary\/day_","{","}","_","{","}",".npy","''","datapipe","=","InMemoryBinaryCriteoIterDataPipe","(","dense_paths=","[","template.format","(","0",",","``","dense","''",")",",","template.format","(","1",",","``","dense","''",")","]",",","sparse_paths=","[","template.format","(","0",",","``","sparse","''",")",",","template.format","(","1",",","``","sparse","''",")","]",",","labels_paths=","[","template.format","(","0",",","``","labels","''",")",",","template.format","(","1",",","``","labels","''",")","]",",","batch_size=1024",",","rank=torch.distributed.get_rank","(",")",",","world_size=torch.distributed.get_world_size","(",")",",",")","batch","=","next","(","iter","(","datapipe",")",")"]
code: def __iter__(self) -> Iterator[Batch]:
      # Invariant: buffer never contains more than batch_size rows.
      buffer: Optional[List[np.ndarray]] = None
  
      def append_to_buffer(
          dense: np.ndarray, sparse: np.ndarray, labels: np.ndarray
      ) -> None:
          nonlocal buffer
          if buffer is None:
              buffer = [dense, sparse, labels]
          else:
              for idx, arr in enumerate([dense, sparse, labels]):
                  buffer[idx] = np.concatenate((buffer[idx], arr))
  
      # Maintain a buffer that can contain up to batch_size rows. Fill buffer as
      # much as possible on each iteration. Only return a new batch when batch_size
      # rows are filled.
      file_idx = 0
      row_idx = 0
      batch_idx = 0
      buffer_row_count = 0
      cur_batch_size = (
          self.batch_size
          if self.num_full_batches > 0
          else self.last_batch_sizes[0]
      )
      while (
          batch_idx
          < self.num_full_batches
          + (self.last_batch_sizes[0] > 0) * self.world_size
      ):
          if buffer_row_count == cur_batch_size or file_idx == len(
              self.dense_arrs
          ):
              if batch_idx % self.world_size == self.rank:
                  yield self._np_arrays_to_batch(*none_throws(buffer))
                  buffer = None
              buffer_row_count = 0
              batch_idx += 1
              if (
                  0
                  <= batch_idx - self.num_full_batches
                  < self.world_size
                  and (self.last_batch_sizes[0] > 0)
              ):
                  cur_batch_size = self.last_batch_sizes[
                      batch_idx - self.num_full_batches
                  ]
          else:
              rows_to_get = min(
                  cur_batch_size - buffer_row_count,
                  self.num_rows_per_file[file_idx] - row_idx,
              )
              buffer_row_count += rows_to_get
              slice_ = slice(row_idx, row_idx + rows_to_get)
  
              if batch_idx % self.world_size == self.rank:
                  dense_inputs = self.dense_arrs[file_idx][slice_, :]
                  sparse_inputs = self.sparse_arrs[file_idx][slice_, :]
                  target_labels = self.labels_arrs[file_idx][slice_, :]
  
                  if self.mmap_mode and self.hashes is not None:
                      sparse_inputs = sparse_inputs % self.hashes
  
                  append_to_buffer(
                      dense_inputs,
                      sparse_inputs,
                      target_labels,
                  )
              row_idx += rows_to_get
  
              if row_idx >= self.num_rows_per_file[file_idx]:
                  file_idx += 1
                  row_idx = 0
  
imports: import os
  import shutil
  import time
  from typing import Any, Callable, Dict, Iterable, Iterator, List, Optional, Tuple, Union
  import numpy
  import torch
  import torch.utils.data.datapipes
  from iopath.common.file_io import PathManager, PathManagerFactory
  from pyre_extensions import none_throws
  from torch.utils.data import IterableDataset, IterDataPipe
  from torchrec.datasets.utils import Batch, LoadFiles, PATH_MANAGER_KEY, ReadLinesFromCSV, safe_cast
  from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
label_desc: Use image node_id 5 for calling the InMemoryBinaryCriteoIterDataPipe obj's underlying member method code with example usage: obj.__iter__() without return types