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: Registers a layer for sparsification. The layer should be part of self.model.
  Specifically, registers a pre-forward hook to the layer. The hook will apply the aggregate_fn
  and store the aggregated activations that is input over each step.
  
  Note::
      - There is no need to pass in the name of the layer as it is automatically computed as per
        the fqn convention.
  
      - All the functions (fn) passed as argument will be called at a dim, feature level.
docstr_tok: ["Registers","a","layer","for","sparsification",".","The","layer","should","be","part","of","self.model",".","Specifically",",","registers","a","pre-forward","hook","to","the","layer",".","The","hook","will","apply","the","aggregate_fn","and","store","the","aggregated","activations","that","is","input","over","each","step",".","Note",":",":","-","There","is","no","need","to","pass","in","the","name","of","the","layer","as","it","is","automatically","computed","as","per","the","fqn","convention",".","-","All","the","functions","(","fn",")","passed","as","argument","will","be","called","at","a","dim",",","feature","level","."]
code: def register_layer(
      self,
      layer: nn.Module,
      aggregate_fn=None,
      reduce_fn=None,
      mask_fn=None,
      features=None,
      feature_dim=None,
      **sparse_config,
  ):
      r"""
      Registers a layer for sparsification. The layer should be part of self.model.
      Specifically, registers a pre-forward hook to the layer. The hook will apply the aggregate_fn
      and store the aggregated activations that is input over each step.
  
      Note::
          - There is no need to pass in the name of the layer as it is automatically computed as per
            the fqn convention.
  
          - All the functions (fn) passed as argument will be called at a dim, feature level.
      """
      name = module_to_fqn(self.model, layer)
      assert (
          name is not None
      ), "layer not found in the model"  # satisfy mypy
  
      if (
          name in self.data_groups
      ):  # unregister layer if already present
          warnings.warn(
              "layer already attached to the sparsifier, deregistering the layer and registering with new config"
          )
          self.unregister_layer(name=name)
  
      local_args = copy.deepcopy(self.defaults)
      update_dict = {
          "aggregate_fn": aggregate_fn,
          "reduce_fn": reduce_fn,
          "mask_fn": mask_fn,
          "features": features,
          "feature_dim": feature_dim,
          "layer": layer,
      }
      local_args.update(
          (arg, val)
          for arg, val in update_dict.items()
          if val is not None
      )
      local_args["sparse_config"].update(sparse_config)
  
      self._safe_rail_checks(local_args)
  
      self.data_groups[name] = local_args
      agg_hook = layer.register_forward_pre_hook(
          self._aggregate_hook(name=name)
      )
  
      self.state[name][
          "mask"
      ] = None  # mask will be created when model forward is called.
  
      # attach agg hook
      self.data_groups[name]["hook"] = agg_hook
  
      # for serialization purposes, we know whether aggregate_hook is attached
      # or sparsify_hook()
      self.data_groups[name]["hook_state"] = "aggregate"
  
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 4 for calling the ActivationSparsifier obj's underlying member method code with example usage: obj.register_layer(layer, aggregate_fn, reduce_fn, mask_fn, features, feature_dim) without return types