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()
code: def __init__(
      self,
      model: nn.Module,
      aggregate_fn=None,
      reduce_fn=None,
      mask_fn=None,
      features=None,
      feature_dim=None,
      **sparse_config,
  ):
      self.model = model
      self.defaults: Dict[str, Any] = defaultdict()
      self.defaults["sparse_config"] = sparse_config
  
      # functions
      self.defaults["aggregate_fn"] = aggregate_fn
      self.defaults["reduce_fn"] = reduce_fn
      self.defaults["mask_fn"] = mask_fn
  
      # default feature and feature_dim
      self.defaults["features"] = features
      self.defaults["feature_dim"] = feature_dim
  
      self.data_groups: Dict[str, Dict] = defaultdict(
          dict
      )  # contains all relevant info w.r.t each registered layer
  
      self.state: Dict[str, Any] = defaultdict(dict)
  
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 1 to create a new ActivationSparsifier object  with example: obj = ActivationSparsifier(model, aggregate_fn, reduce_fn, mask_fn, features, feature_dim)