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: Args:
      m: nn.Module to be wrapped
      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"
      strict: If set to ``True``, the input module will be validated to check that
          ``GradSampleModule`` has grad sampler functions for all submodules of
          the input module (i.e. if it knows how to calculate per sample gradients)
          for all model parameters. If set to ``False``, per sample gradients will
          be computed on "best effort" basis - they will be available where
          possible and set to None otherwise. This is not recommended, because
          some unsupported modules (e.g. BatchNorm) affect other parameters and
          invalidate the concept of per sample gradients for the entire model.
      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.
  
  Raises:
      NotImplementedError
          If ``strict`` is set to ``True`` and module ``m`` (or any of its
          submodules) doesn't have a registered grad sampler function.
docstr_tok: ["Args",":","m",":","nn.Module","to","be","wrapped","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","''","strict",":","If","set","to","``","True","``",",","the","input","module","will","be","validated","to","check","that","``","GradSampleModule","``","has","grad","sampler","functions","for","all","submodules","of","the","input","module","(","i.e",".","if","it","knows","how","to","calculate","per","sample","gradients",")","for","all","model","parameters",".","If","set","to","``","False","``",",","per","sample","gradients","will","be","computed","on","``","best","effort","''","basis","-","they","will","be","available","where","possible","and","set","to","None","otherwise",".","This","is","not","recommended",",","because","some","unsupported","modules","(","e.g",".","BatchNorm",")","affect","other","parameters","and","invalidate","the","concept","of","per","sample","gradients","for","the","entire","model",".","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",".","Raises",":","NotImplementedError","If","``","strict","``","is","set","to","``","True","``","and","module","``","m","``","(","or","any","of","its","submodules",")","does","n't","have","a","registered","grad","sampler","function","."]
code: def __init__(
      self,
      m: nn.Module,
      *,
      batch_first=True,
      loss_reduction="mean",
      strict: bool = True,
      force_functorch=False,
  ):
      """
  
      Args:
          m: nn.Module to be wrapped
          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"
          strict: If set to ``True``, the input module will be validated to check that
              ``GradSampleModule`` has grad sampler functions for all submodules of
              the input module (i.e. if it knows how to calculate per sample gradients)
              for all model parameters. If set to ``False``, per sample gradients will
              be computed on "best effort" basis - they will be available where
              possible and set to None otherwise. This is not recommended, because
              some unsupported modules (e.g. BatchNorm) affect other parameters and
              invalidate the concept of per sample gradients for the entire model.
          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.
  
      Raises:
          NotImplementedError
              If ``strict`` is set to ``True`` and module ``m`` (or any of its
              submodules) doesn't have a registered grad sampler function.
      """
      super().__init__(
          m,
          batch_first=batch_first,
          loss_reduction=loss_reduction,
      )
  
      errors = self.validate(module=m, strict=strict)
      if errors and not strict:
          logger.info(
              f"GradSampleModule found the following errors: {errors}."
              "Using non-strict mode, continuing"
          )
  
      self.hooks_enabled = False
      self.grad_accumulation_allowed = True
      self.batch_first = batch_first
      self.loss_reduction = loss_reduction
      self.force_functorch = force_functorch
      self.add_hooks(
          loss_reduction=loss_reduction,
          batch_first=batch_first,
          force_functorch=force_functorch,
      )
  
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 1 to create a new GradSampleModule object from inherited base classes: AbstractGradSampleModule with example: obj = GradSampleModule(m)