class_docstr: Distributed batch sampler.
  
  Each batch is sampled as follows:
      1. Shuffle the dataset (enabled by default)
      2. Split the dataset among the replicas into chunks of equal size
         (plus or minus one sample)
      3. Each replica selects each sample of its chunk independently
         with probability `sample_rate`
      4. Each replica outputs the selected samples, which form a local batch
  
  The sum of the lengths of the local batches follows a Poisson distribution.
  In particular, the expected length of each local batch is:
  `sample_rate * total_size / num_replicas`
class_docstr_tok: ["Distributed","batch","sampler",".","Each","batch","is","sampled","as","follows",":","1",".","Shuffle","the","dataset","(","enabled","by","default",")","2",".","Split","the","dataset","among","the","replicas","into","chunks","of","equal","size","(","plus","or","minus","one","sample",")","3",".","Each","replica","selects","each","sample","of","its","chunk","independently","with","probability","`","sample_rate","`","4",".","Each","replica","outputs","the","selected","samples",",","which","form","a","local","batch","The","sum","of","the","lengths","of","the","local","batches","follows","a","Poisson","distribution",".","In","particular",",","the","expected","length","of","each","local","batch","is",":","`","sample_rate","*","total_size","\/","num_replicas","`"]
code: def __iter__(self):
      if self.shuffle:
          # deterministically shuffle based on epoch and seed
          g = torch.Generator()
          g.manual_seed(self.shuffle_seed + self.epoch)
          indices = torch.randperm(self.total_size, generator=g)  # type: ignore
      else:
          indices = torch.arange(self.total_size)  # type: ignore
  
      # Subset of the dataset assigned to this replica
      # NOTE: the first replicas might have 1 more sample.
      # (Different from the regular distributed loader that pads with more samples)
      indices = indices[self.rank : self.total_size : self.num_replicas]
      assert len(indices) == self.num_samples
  
      # Now, select a batch with Poisson subsampling
      for _ in range(self.num_batches):
          mask = (
              torch.rand(self.num_samples, generator=self.generator)
              < self.sample_rate
          )
          selected_examples = mask.nonzero(as_tuple=False).reshape(-1)
          if len(selected_examples) > 0:
              yield indices[selected_examples]
  
file_path: opacus/opacus/utils/uniform_sampler.py
imports: from typing import List
  import torch
  from torch.utils.data import Sampler
label_desc: Use image node_id 2 for calling the DistributedUniformWithReplacementSampler obj's underlying member method code with example usage: obj.__iter__() without return types