code: def compute_microbatch_grad_sample(
      x: Union[torch.Tensor, List[torch.Tensor]],
      module: nn.Module,
      batch_first: bool = True,
      loss_reduction: str = "mean",
      chunk_method: Callable = iter,
  ) -> Dict[str, torch.tensor]:
      """
      Computes per-sample gradients with the microbatch method, i.e. by computing normal gradients
      with batch_size set to 1, and manually accumulating them. This is our reference for testing
      as this method is obviously correct, but slow.
  
      Args:
          x: Sample input batch
          module: The nn.Module you want to test.
          batch_first: Whether batch size is the first dimension (as opposed to the second).
              Defaults to True.
          loss_reduction: What reduction to apply to the loss. Defaults to "mean".
          chunk_method: The method to use to split the batch into microbatches. Defaults to ``iter``.
  
      Returns:
          Dictionary mapping parameter_name -> per-sample-gradient for that parameter
      """
      torch.use_deterministic_algorithms(True)
      torch.manual_seed(0)
      np.random.seed(0)
  
      module = ModelWithLoss(clone_module(module), loss_reduction)
  
      for _, p in trainable_parameters(module):
          p.microbatch_grad_sample = []
  
      if not batch_first and type(x) is not list:
          # This allows us to iterate with x_i
          x = x.transpose(0, 1)
  
      # Invariant: x is [B, T, ...]
  
      for x_i in chunk_method(x):
          # x_i is [T, ...]
          module.zero_grad()
          if type(x_i) is not tuple:
              # EmbeddingBag provides tuples
              x_i = x_i.unsqueeze(
                  0 if batch_first else 1
              )  # x_i of size [1, T, ...] if batch_first, else [T, 1, ...]
          loss_i = module(x_i)
          loss_i.backward()
          for p in module.parameters():
              p.microbatch_grad_sample.append(p.grad.detach().clone())
  
      for _, p in trainable_parameters(module):
          if batch_first:
              p.microbatch_grad_sample = torch.stack(
                  p.microbatch_grad_sample, dim=0  # [B, T, ...]
              )
          else:
              p.microbatch_grad_sample = torch.stack(
                  p.microbatch_grad_sample, dim=1  # [T, B, ...]
              ).transpose(
                  0, 1
              )  # Opacus's semantics is that grad_samples are ALWAYS batch_first: [B, T, ...]
  
      microbatch_grad_samples = {
          name: p.microbatch_grad_sample
          for name, p in trainable_parameters(module.wrapped_module)
      }
      return microbatch_grad_samples
  
imports: import io
  from typing import Callable, Dict, Iterable, List, Union
  import numpy
  import torch
  import torch.nn
  from opacus.grad_sample import wrap_model
  from opacus.utils.module_utils import trainable_parameters
  from opacus.utils.packed_sequences import compute_seq_lengths
  from torch.nn.utils.rnn import PackedSequence, pad_packed_sequence
label_desc: Use image node_id 3 for calling a global function with example usage: compute_microbatch_grad_sample(x, module, batch_first, loss_reduction, chunk_method) and returns: microbatch_grad_samples