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 _load_data_for_rank(self) -> None:
      start_row, last_row = 0, None
      if self.stage in ["val", "test"]:
          # Last day's dataset is split into 2 sets: 1st half for "val"; 2nd for "test"
          samples_in_file = BinaryCriteoUtils.get_shape_from_npy(
              self.dense_paths[0],
              path_manager_key=self.path_manager_key,
          )[0]
          start_row = 0
          dataset_len = int(np.ceil(samples_in_file / 2.0))
          if self.stage == "test":
              start_row = dataset_len
              dataset_len = samples_in_file - start_row
          last_row = start_row + dataset_len - 1
  
      (
          row_ranges,
          remainder,
      ) = BinaryCriteoUtils.get_file_row_ranges_and_remainder(
          lengths=[
              BinaryCriteoUtils.get_shape_from_npy(
                  path, path_manager_key=self.path_manager_key
              )[0]
              for path in self.dense_paths
          ],
          rank=self.rank,
          world_size=self.world_size,
          start_row=start_row,
          last_row=last_row,
      )
      self.remainder = remainder
      self.dense_arrs, self.sparse_arrs, self.labels_arrs = [], [], []
      for arrs, paths in zip(
          [self.dense_arrs, self.sparse_arrs, self.labels_arrs],
          [self.dense_paths, self.sparse_paths, self.labels_paths],
      ):
          for idx, (range_left, range_right) in row_ranges.items():
              arrs.append(
                  BinaryCriteoUtils.load_npy_range(
                      paths[idx],
                      range_left,
                      range_right - range_left + 1,
                      path_manager_key=self.path_manager_key,
                      mmap_mode=self.mmap_mode,
                  )
              )
  
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 2 for calling the InMemoryBinaryCriteoIterDataPipe obj's underlying member method code with example usage: obj._load_data_for_rank() without return types