class_docstr: The Activation sparsifier class aims to sparsify/prune activations in a neural
  network. The idea is to attach the sparsifier to a layer (or layers) and it
  zeroes out the activations based on the mask_fn (or sparsification function)
  input by the user.
  The mask_fn is applied once all the inputs are aggregated and reduced i.e.
  mask = mask_fn(reduce_fn(aggregate_fn(activations)))
  
  Note::
      The sparsification mask is computed on the input **before it goes through the attached layer**.
  
  Args:
      model (nn.Module):
          The model whose layers will be sparsified. The layers that needs to be
          sparsified should be added separately using the register_layer() function
      aggregate_fn (Optional, Callable):
          default aggregate_fn that is used if not specified while registering the layer.
          specifies how inputs should be aggregated over time.
          The aggregate_fn should usually take 2 torch tensors and return the aggregated tensor.
          Example
              def add_agg_fn(tensor1, tensor2):  return tensor1 + tensor2
              reduce_fn (Optional, Callable):
                  default reduce_fn that is used if not specified while registering the layer.
                  reduce_fn will be called on the aggregated tensor i.e. the tensor obtained after
                  calling agg_fn() on all inputs.
                  Example
              def mean_reduce_fn(agg_tensor):    return agg_tensor.mean(dim=0)
              mask_fn (Optional, Callable):
                  default mask_fn that is used to create the sparsification mask using the tensor obtained after
                  calling the reduce_fn(). This is used by default if a custom one is passed in the
                  register_layer().
                  Note that the mask_fn() definition should contain the sparse arguments that is passed in sparse_config
                  arguments.
              features (Optional, list):
                  default selected features to sparsify.
                  If this is non-empty, then the mask_fn will be applied for each feature of the input.
                  For example,
              mask = [mask_fn(reduce_fn(aggregated_fn(input[feature])) for feature in features]
              feature_dim (Optional, int):
                  default dimension of input features. Again, features along this dim will be chosen
                  for sparsification.
              sparse_config (Dict):
                  Default configuration for the mask_fn. This config will be passed
                  with the mask_fn()
  
  Example:
      >>> # xdoctest: +SKIP
      >>> model = SomeModel()
      >>> act_sparsifier = ActivationSparsifier(...)  # init activation sparsifier
      >>> # Initialize aggregate_fn
      >>> def agg_fn(x, y):
      >>>     return x + y
      >>>
      >>> # Initialize reduce_fn
      >>> def reduce_fn(x):
      >>>     return torch.mean(x, dim=0)
      >>>
      >>> # Initialize mask_fn
      >>> def mask_fn(data):
      >>>     return torch.eye(data.shape).to(data.device)
      >>>
      >>>
      >>> act_sparsifier.register_layer(model.some_layer, aggregate_fn=agg_fn, reduce_fn=reduce_fn, mask_fn=mask_fn)
      >>>
      >>> # start training process
      >>> for _ in [...]:
      >>>     # epoch starts
      >>>         # model.forward(), compute_loss() and model.backwards()
      >>>     # epoch ends
      >>>     act_sparsifier.step()
      >>> # end training process
      >>> sparsifier.squash_mask()
docstr: Returns mask associated to the layer.
  
  The mask is
      - a torch tensor is features for that layer is None.
      - a list of torch tensors for each feature, otherwise
  
  Note::
      The shape of the mask is unknown until model.forward() is applied.
      Hence, if get_mask() is called before model.forward(), an
      error will be raised.
docstr_tok: ["Returns","mask","associated","to","the","layer",".","The","mask","is","-","a","torch","tensor","is","features","for","that","layer","is","None",".","-","a","list","of","torch","tensors","for","each","feature",",","otherwise","Note",":",":","The","shape","of","the","mask","is","unknown","until","model.forward","(",")","is","applied",".","Hence",",","if","get_mask","(",")","is","called","before","model.forward","(",")",",","an","error","will","be","raised","."]
code: def get_mask(
      self,
      name: Optional[str] = None,
      layer: Optional[nn.Module] = None,
  ):
      """
      Returns mask associated to the layer.
  
      The mask is
          - a torch tensor is features for that layer is None.
          - a list of torch tensors for each feature, otherwise
  
      Note::
          The shape of the mask is unknown until model.forward() is applied.
          Hence, if get_mask() is called before model.forward(), an
          error will be raised.
      """
      assert (
          name is not None or layer is not None
      ), "Need at least name or layer obj to retrieve mask"
  
      if name is None:
          assert layer is not None
          name = module_to_fqn(self.model, layer)
          assert (
              name is not None
          ), "layer not found in the specified model"
  
      if name not in self.state:
          raise ValueError("Error: layer with the given name not found")
  
      mask = self.state[name].get("mask", None)
  
      if mask is None:
          raise ValueError(
              "Error: shape unknown, call layer() routine at least once to infer mask"
          )
      return mask
  
imports: from typing import Any, Dict, List, Optional
  import torch
  from collections import defaultdict
  from torch import nn
  import copy
  from ...sparsifier.utils import fqn_to_module, module_to_fqn
  import warnings
label_desc: Use image node_id 5 for calling the ActivationSparsifier obj's underlying member method code with example usage: obj.get_mask(name, layer) and returns: mask