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: Computes per sample gradients given the current backprops and activations
  stored by the associated forward hook. Computed per sample gradients are
  stored in ``grad_sample`` field in each parameter.
  
  For non-recurrent layers the process is straightforward: for each
  ``loss.backward()`` call this hook will be called exactly one. For recurrent
  layers, however, this is more complicated and the hook will be called multiple
  times, while still processing the same batch of data.
  
  For this reason we first accumulate the gradients from *the same batch* in
  ``p._current_grad_sample`` and then, when we detect the end of a full backward
  pass - we store accumulated result on ``p.grad_sample``.
  
  From there, ``p.grad_sample`` could be either a Tensor or a list of Tensors,
  if accumulated over multiple batches
  
  Args:
      module: nn.Module,
      _forward_input: torch.Tensor,
      forward_output: torch.Tensor,
      loss_reduction: str,
      batch_first: bool,
docstr_tok: ["Computes","per","sample","gradients","given","the","current","backprops","and","activations","stored","by","the","associated","forward","hook",".","Computed","per","sample","gradients","are","stored","in","``","grad_sample","``","field","in","each","parameter",".","For","non-recurrent","layers","the","process","is","straightforward",":","for","each","``","loss.backward","(",")","``","call","this","hook","will","be","called","exactly","one",".","For","recurrent","layers",",","however",",","this","is","more","complicated","and","the","hook","will","be","called","multiple","times",",","while","still","processing","the","same","batch","of","data",".","For","this","reason","we","first","accumulate","the","gradients","from","*","the","same","batch","*","in","``","p._current_grad_sample","``","and","then",",","when","we","detect","the","end","of","a","full","backward","pass","-","we","store","accumulated","result","on","``","p.grad_sample","``",".","From","there",",","``","p.grad_sample","``","could","be","either","a","Tensor","or","a","list","of","Tensors",",","if","accumulated","over","multiple","batches","Args",":","module",":","nn.Module",",","_forward_input",":","torch.Tensor",",","forward_output",":","torch.Tensor",",","loss_reduction",":","str",",","batch_first",":","bool",","]
code: def capture_backprops_hook(
      self,
      module: nn.Module,
      _forward_input: torch.Tensor,
      forward_output: torch.Tensor,
      loss_reduction: str,
      batch_first: bool,
  ):
      """
      Computes per sample gradients given the current backprops and activations
      stored by the associated forward hook. Computed per sample gradients are
      stored in ``grad_sample`` field in each parameter.
  
      For non-recurrent layers the process is straightforward: for each
      ``loss.backward()`` call this hook will be called exactly one. For recurrent
      layers, however, this is more complicated and the hook will be called multiple
      times, while still processing the same batch of data.
  
      For this reason we first accumulate the gradients from *the same batch* in
      ``p._current_grad_sample`` and then, when we detect the end of a full backward
      pass - we store accumulated result on ``p.grad_sample``.
  
      From there, ``p.grad_sample`` could be either a Tensor or a list of Tensors,
      if accumulated over multiple batches
  
      Args:
          module: nn.Module,
          _forward_input: torch.Tensor,
          forward_output: torch.Tensor,
          loss_reduction: str,
          batch_first: bool,
      """
      if not self.hooks_enabled:
          return
  
      backprops = forward_output[0].detach()
      activations, backprops = self.rearrange_grad_samples(
          module=module,
          backprops=backprops,
          loss_reduction=loss_reduction,
          batch_first=batch_first,
      )
      if (
          not self.force_functorch
          and type(module) in self.GRAD_SAMPLERS
      ):
          grad_sampler_fn = self.GRAD_SAMPLERS[type(module)]
      else:
          grad_sampler_fn = ft_compute_per_sample_gradient
  
      grad_samples = grad_sampler_fn(module, activations, backprops)
      for param, gs in grad_samples.items():
          create_or_accumulate_grad_sample(
              param=param,
              grad_sample=gs,
              max_batch_len=module.max_batch_len,
          )
  
      # Detect end of current batch processing and switch accumulation
      # mode from sum to stacking. Used for RNNs and tied parameters
      # (See #417 for details)
      for _, p in trainable_parameters(module):
          p._forward_counter -= 1
          if p._forward_counter == 0:
              promote_current_grad_sample(p)
  
          if not self.grad_accumulation_allowed:
              if (
                  isinstance(p.grad_sample, list)
                  and len(p.grad_sample) > 1
              ):
                  raise ValueError(
                      "Poisson sampling is not compatible with grad accumulation. "
                      "You need to call optimizer.step() after every forward/backward pass "
                      "or consider using BatchMemoryManager"
                  )
  
      if len(module.activations) == 0:
          if hasattr(module, "max_batch_len"):
              del module.max_batch_len
  
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 10 for calling the GradSampleModule obj's underlying member method code with example usage: obj.capture_backprops_hook(module, _forward_input, forward_output, loss_reduction, batch_first) without return types