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: Adds hooks to model to save activations and backprop values.
  The hooks will
  1. save activations into param.activations during forward pass
  2. compute per-sample gradients in params.grad_sample during backward pass.
  Call ``remove_hooks(model)`` to disable this.
  
  Args:
      model: the model to which hooks are added
      batch_first: Flag to indicate if the input tensor to the corresponding module
          has the first dimension representing the batch. If set to True, dimensions on
          input tensor are expected be ``[batch_size, ...]``, otherwise
          ``[K, batch_size, ...]``
      loss_reduction: Indicates if the loss reduction (for aggregating the gradients)
          is a sum or a mean operation. Can take values "sum" or "mean"
      force_functorch: If set to ``True``, will use functorch to compute all per sample gradients.
          Otherwise, functorch will be used only for layers without registered grad sampler methods.
docstr_tok: ["Adds","hooks","to","model","to","save","activations","and","backprop","values",".","The","hooks","will","1.","save","activations","into","param.activations","during","forward","pass","2.","compute","per-sample","gradients","in","params.grad_sample","during","backward","pass",".","Call","``","remove_hooks","(","model",")","``","to","disable","this",".","Args",":","model",":","the","model","to","which","hooks","are","added","batch_first",":","Flag","to","indicate","if","the","input","tensor","to","the","corresponding","module","has","the","first","dimension","representing","the","batch",".","If","set","to","True",",","dimensions","on","input","tensor","are","expected","be","``","[","batch_size",",","...","]","``",",","otherwise","``","[","K",",","batch_size",",","...","]","``","loss_reduction",":","Indicates","if","the","loss","reduction","(","for","aggregating","the","gradients",")","is","a","sum","or","a","mean","operation",".","Can","take","values","``","sum","''","or","``","mean","''","force_functorch",":","If","set","to","``","True","``",",","will","use","functorch","to","compute","all","per","sample","gradients",".","Otherwise",",","functorch","will","be","used","only","for","layers","without","registered","grad","sampler","methods","."]
code: def add_hooks(
      self,
      *,
      loss_reduction: str = "mean",
      batch_first: bool = True,
      force_functorch: bool = False,
  ) -> None:
      """
      Adds hooks to model to save activations and backprop values.
      The hooks will
      1. save activations into param.activations during forward pass
      2. compute per-sample gradients in params.grad_sample during backward pass.
      Call ``remove_hooks(model)`` to disable this.
  
      Args:
          model: the model to which hooks are added
          batch_first: Flag to indicate if the input tensor to the corresponding module
              has the first dimension representing the batch. If set to True, dimensions on
              input tensor are expected be ``[batch_size, ...]``, otherwise
              ``[K, batch_size, ...]``
          loss_reduction: Indicates if the loss reduction (for aggregating the gradients)
              is a sum or a mean operation. Can take values "sum" or "mean"
          force_functorch: If set to ``True``, will use functorch to compute all per sample gradients.
              Otherwise, functorch will be used only for layers without registered grad sampler methods.
      """
      if hasattr(self._module, "autograd_grad_sample_hooks"):
          raise ValueError(
              "Trying to add hooks twice to the same model"
          )
      else:
          self._module.autograd_grad_sample_hooks = []
          self.autograd_grad_sample_hooks = (
              self._module.autograd_grad_sample_hooks
          )
  
      for module in self.iterate_submodules(self._module):
          # Do not add hooks to DPRNN, DPLSTM or DPGRU as the hooks are handled by the `RNNLinear`
          if type(module) in [DPRNN, DPLSTM, DPGRU]:
              continue
  
          if force_functorch or not type(module) in self.GRAD_SAMPLERS:
              prepare_layer(module, batch_first=batch_first)
  
          self.autograd_grad_sample_hooks.append(
              module.register_forward_hook(
                  self.capture_activations_hook
              )
          )
  
          self.autograd_grad_sample_hooks.append(
              module.register_backward_hook(
                  partial(
                      self.capture_backprops_hook,
                      loss_reduction=loss_reduction,
                      batch_first=batch_first,
                  )
              )
          )
  
      self.enable_hooks()
  
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 4 for calling the GradSampleModule obj's underlying member method code with example usage: obj.add_hooks() without return types