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 __init__(
      self,
      stage: str,
      dense_paths: List[str],
      sparse_paths: List[str],
      labels_paths: List[str],
      batch_size: int,
      rank: int,
      world_size: int,
      drop_last: Optional[bool] = False,
      shuffle_batches: bool = False,
      shuffle_training_set: bool = False,
      shuffle_training_set_random_seed: int = 0,
      mmap_mode: bool = False,
      hashes: Optional[List[int]] = None,
      path_manager_key: str = PATH_MANAGER_KEY,
  ) -> None:
      self.stage = stage
      self.dense_paths = dense_paths
      self.sparse_paths = sparse_paths
      self.labels_paths = labels_paths
      self.batch_size = batch_size
      self.rank = rank
      self.world_size = world_size
      self.drop_last = drop_last
      self.shuffle_batches = shuffle_batches
      self.shuffle_training_set = shuffle_training_set
      np.random.seed(shuffle_training_set_random_seed)
      self.mmap_mode = mmap_mode
      self.hashes: np.ndarray = np.array(hashes).reshape(
          (1, CAT_FEATURE_COUNT)
      )
      self.path_manager_key = path_manager_key
      self.path_manager: PathManager = PathManagerFactory().get(
          path_manager_key
      )
  
      if shuffle_training_set and stage == "train":
          self._shuffle_and_load_data_for_rank()
          self.world_size = 1
          self.rank = 0
      else:
          m = "r" if mmap_mode else None
          self.dense_arrs: List[np.ndarray] = [
              np.load(f, mmap_mode=m) for f in self.dense_paths
          ]
          self.sparse_arrs: List[np.ndarray] = [
              np.load(f, mmap_mode=m) for f in self.sparse_paths
          ]
          self.labels_arrs: List[np.ndarray] = [
              np.load(f, mmap_mode=m) for f in self.labels_paths
          ]
      len_d0 = len(self.dense_arrs[0])
      second_half_start_index = int(len_d0 // 2 + len_d0 % 2)
      if stage == "val":
          self.dense_arrs[0] = self.dense_arrs[0][
              :second_half_start_index, :
          ]
          self.sparse_arrs[0] = self.sparse_arrs[0][
              :second_half_start_index, :
          ]
          self.labels_arrs[0] = self.labels_arrs[0][
              :second_half_start_index, :
          ]
      elif stage == "test":
          self.dense_arrs[0] = self.dense_arrs[0][
              second_half_start_index:, :
          ]
          self.sparse_arrs[0] = self.sparse_arrs[0][
              second_half_start_index:, :
          ]
          self.labels_arrs[0] = self.labels_arrs[0][
              second_half_start_index:, :
          ]
      # When mmap_mode is enabled, sparse features are hashed when
      # samples are batched in def __iter__. Otherwise, the dataset has been
      # preloaded with sparse features hashed in the preload stage, here:
      if not self.mmap_mode and self.hashes is not None:
          for sparse_arr in self.sparse_arrs:
              sparse_arr %= self.hashes
  
      self.num_rows_per_file: List[int] = list(
          map(len, self.dense_arrs)
      )
      total_rows = sum(self.num_rows_per_file)
      self.num_full_batches: int = (
          total_rows // batch_size // self.world_size * self.world_size
      )
      self.last_batch_sizes: np.ndarray = np.array(
          [0 for _ in range(self.world_size)]
      )
      remainder = total_rows % (self.world_size * batch_size)
      if not self.drop_last and 0 < remainder:
          if remainder < self.world_size:
              self.num_full_batches -= self.world_size
              self.last_batch_sizes += batch_size
          else:
              self.last_batch_sizes += remainder // self.world_size
          self.last_batch_sizes[: remainder % self.world_size] += 1
  
      # These values are the same for the KeyedJaggedTensors in all batches, so they
      # are computed once here. This avoids extra work from the KeyedJaggedTensor sync
      # functions.
      self._num_ids_in_batch: int = CAT_FEATURE_COUNT * (batch_size + 1)
      self.keys: List[str] = DEFAULT_CAT_NAMES
      self.lengths: torch.Tensor = torch.ones(
          (self._num_ids_in_batch,), dtype=torch.int32
      )
      self.offsets: torch.Tensor = torch.arange(
          0, self._num_ids_in_batch + 1, dtype=torch.int32
      )
      self._num_ids_in_batch -= CAT_FEATURE_COUNT
      self.length_per_key: List[int] = CAT_FEATURE_COUNT * [batch_size]
      self.offset_per_key: List[int] = [
          batch_size * i for i in range(CAT_FEATURE_COUNT + 1)
      ]
      self.index_per_key: Dict[str, int] = {
          key: i for (i, key) in enumerate(self.keys)
      }
  
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 1 to create a new InMemoryBinaryCriteoIterDataPipe object from inherited base classes: IterableDataset with example: obj = InMemoryBinaryCriteoIterDataPipe(stage, dense_paths, sparse_paths, labels_paths, batch_size, rank, world_size, drop_last, shuffle_batches, shuffle_training_set, shuffle_training_set_random_seed, mmap_mode, hashes, path_manager_key)