class_docstr: Hooks-based implementation of AbstractGradSampleModule
  
  Computes per-sample gradients using custom-written methods for each layer.
  See README.md for more details
docstr: Rearrange activations and grad_samples based on loss reduction and batch dim
  
  Args:
      module: the module for which per-sample gradients are computed
      backprops: the captured backprops
      loss_reduction: either "mean" or "sum" depending on whether backpropped
          loss was averaged or summed over batch
      batch_first: True is batch dimension is first
docstr_tok: ["Rearrange","activations","and","grad_samples","based","on","loss","reduction","and","batch","dim","Args",":","module",":","the","module","for","which","per-sample","gradients","are","computed","backprops",":","the","captured","backprops","loss_reduction",":","either","``","mean","''","or","``","sum","''","depending","on","whether","backpropped","loss","was","averaged","or","summed","over","batch","batch_first",":","True","is","batch","dimension","is","first"]
code: def rearrange_grad_samples(
      self,
      *,
      module: nn.Module,
      backprops: torch.Tensor,
      loss_reduction: str,
      batch_first: bool,
  ) -> Tuple[torch.Tensor, torch.Tensor]:
      """
      Rearrange activations and grad_samples based on loss reduction and batch dim
  
      Args:
          module: the module for which per-sample gradients are computed
          backprops: the captured backprops
          loss_reduction: either "mean" or "sum" depending on whether backpropped
              loss was averaged or summed over batch
          batch_first: True is batch dimension is first
      """
      if not hasattr(module, "activations"):
          raise ValueError(
              f"No activations detected for {type(module)},"
              " run forward after add_hooks(model)"
          )
  
      batch_dim = 0 if batch_first or type(module) is RNNLinear else 1
  
      if not hasattr(module, "max_batch_len"):
          # For packed sequences, max_batch_len is set in the forward of the model (e.g. the LSTM)
          # Otherwise we infer it here
          module.max_batch_len = _get_batch_size(
              module=module,
              batch_dim=batch_dim,
          )
      activations = module.activations.pop()
  
      n = module.max_batch_len
      if loss_reduction == "mean":
          backprops = backprops * n
      elif loss_reduction == "sum":
          backprops = backprops
      else:
          raise ValueError(
              f"loss_reduction = {loss_reduction}. Only 'sum' and 'mean' losses are supported"
          )
  
      # No matter where the batch dimension was, .grad_samples will *always* put it in the first dim
      if batch_dim != 0:
          activations = [
              t.permute(
                  [batch_dim]
                  + [x for x in range(t.dim()) if x != batch_dim]
              )
              for t in activations
          ]
          backprops = backprops.permute(
              [batch_dim]
              + [x for x in range(backprops.dim()) if x != batch_dim]
          )
  
      return activations, backprops
  
imports: from __future__ import annotations
  import logging
  import warnings
  from functools import partial
  from typing import Iterable, List, Tuple
  import torch
  import torch.nn
  from opacus.grad_sample.functorch import ft_compute_per_sample_gradient, prepare_layer
  from opacus.grad_sample.gsm_base import AbstractGradSampleModule
  from opacus.layers.dp_rnn import DPGRU, DPLSTM, DPRNN, RNNLinear
  from opacus.utils.module_utils import has_trainable_params, requires_grad, trainable_modules, trainable_parameters
label_desc: Use image node_id 11 for calling the GradSampleModule obj's underlying member method code with example usage: obj.rearrange_grad_samples() and returns: activations,  backprops