code: def shuffle(dataset: Iterator[_T], buffer_size: int) -> Iterator[_T]:
      buffer = [next(dataset) for _ in range(buffer_size)]
      random.shuffle(buffer)
      for item in dataset:
          idx = random.randint(0, buffer_size - 1)  # Inclusive.
          result = buffer[idx]
          buffer[idx] = item
          yield result
  
file_path: dm-haiku/examples/transformer/dataset.py
imports: from collections.abc import Iterable, Iterator
  import itertools
  import random
  from typing import NamedTuple, TypeVar
  import numpy
label_desc: Use image node_id 2 for calling a global function with example usage: shuffle(dataset, buffer_size) without return types