code: def _test_in_memory_training_set_shuffle(
      self,
      rows_per_file: List[int],
      batch_size: int,
      world_size: int,
      random_seed: int = 0,
  ) -> None:
      with contextlib.ExitStack() as stack:
          num_rows_csr = np.cumsum([0] + rows_per_file)
          dense, sparse, labels = [], [], []
          for i, _ in enumerate(rows_per_file):
              start = num_rows_csr[i]
              end = num_rows_csr[i + 1]
              dense.append(np.mgrid[start:end, 0:INT_FEATURE_COUNT][0])
              sparse.append(np.mgrid[start:end, 0:CAT_FEATURE_COUNT][0])
              labels.append(np.mgrid[start:end, 0:1][0])
          files = [
              stack.enter_context(
                  self._create_dataset_npys(
                      num_rows=num_rows,
                      dense=dense[i],
                      sparse=sparse[i],
                      labels=labels[i],
                  )
              )
              for i, num_rows in enumerate(rows_per_file)
          ]
          hashes = [i + 1 for i in range(CAT_FEATURE_COUNT)]
          dataset_len = num_rows = sum(rows_per_file)
          remainder = dataset_len % world_size
          rows_per_rank = dataset_len // world_size
          rows_per_rank = np.array(
              [rows_per_rank for _ in range(world_size)]
          )
          rows_per_rank[:remainder] += 1
          lens = []
          total_samples_count = 0
          for rank in range(world_size):
              shuffle_dataset_len = rows_per_rank[rank]
              num_batches = math.ceil(shuffle_dataset_len / batch_size)
              end_batch_size = shuffle_dataset_len % batch_size
  
              datapipe = InMemoryBinaryCriteoIterDataPipe(
                  stage="train",
                  dense_paths=[f[0] for f in files],
                  sparse_paths=[f[1] for f in files],
                  labels_paths=[f[2] for f in files],
                  batch_size=batch_size,
                  rank=rank,
                  world_size=world_size,
                  shuffle_training_set=True,
                  shuffle_training_set_random_seed=random_seed,
                  hashes=hashes,
              )
              datapipe_len = len(datapipe)
              self.assertEqual(datapipe_len, num_batches)
  
              np.random.seed(random_seed)
              permutation_arr = np.random.permutation(num_rows)
              src_arr = np.arange(num_rows)
              target_shuffled_arr = np.empty(num_rows)
              target_shuffled_arr[permutation_arr] = src_arr
              target_shuffled_arr_ = np.array(
                  [b for a, b in sorted(zip(permutation_arr, src_arr))]
              )
              np.testing.assert_array_equal(
                  target_shuffled_arr, target_shuffled_arr_
              )
              len_ = 0
              for batch in datapipe:
                  if len_ < num_batches - 1 or end_batch_size == 0:
                      self._validate_batch(batch, batch_size=batch_size)
                  else:
                      self._validate_batch(
                          batch, batch_size=end_batch_size
                      )
                  len_ += 1
                  total_samples_count += batch.dense_features.shape[0]
  
              # Check that dataset __len__ matches true length.
              self.assertEqual(datapipe_len, len_)
              lens.append(len_)
  
          # Ensure all ranks return the same number of batches.
          self.assertEqual(len(set(lens)), 1)
          # Ensure the number of samples read match the number of samples in the dataset
          self.assertEqual(dataset_len, total_samples_count)
  
imports: import contextlib
  import math
  import os
  import random
  import tempfile
  from typing import Any, Dict, List, Optional
  import numpy
  from torch.utils.data import DataLoader
  from torchrec.datasets.criteo import BinaryCriteoUtils, CAT_FEATURE_COUNT, criteo_kaggle, criteo_terabyte, InMemoryBinaryCriteoIterDataPipe, INT_FEATURE_COUNT
  from torchrec.datasets.test_utils.criteo_test_utils import CriteoTest
  from torchrec.datasets.utils import Batch
label_desc: Use image node_id 3 for calling the TestInMemoryBinaryCriteoIterDataPipe obj's underlying member method code with example usage: obj._test_in_memory_training_set_shuffle(rows_per_file, batch_size, world_size, random_seed) without return types