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: Called for each registered layer and does the following-
      1. apply reduce_fn on the aggregated activations
      2. use mask_fn to compute the sparsification mask
  
  Note:
      the reduce_fn and mask_fn is called for each feature, dim over the data
docstr_tok: ["Called","for","each","registered","layer","and","does","the","following-","1.","apply","reduce_fn","on","the","aggregated","activations","2.","use","mask_fn","to","compute","the","sparsification","mask","Note",":","the","reduce_fn","and","mask_fn","is","called","for","each","feature",",","dim","over","the","data"]
code: def update_mask(self, name, data, configs):
      """
      Called for each registered layer and does the following-
          1. apply reduce_fn on the aggregated activations
          2. use mask_fn to compute the sparsification mask
  
      Note:
          the reduce_fn and mask_fn is called for each feature, dim over the data
      """
      mask = self.get_mask(name)
      sparse_config = configs["sparse_config"]
      features = configs["features"]
      reduce_fn = configs["reduce_fn"]
      mask_fn = configs["mask_fn"]
      if features is None:
          data = reduce_fn(data)
          mask.data = mask_fn(data, **sparse_config)
      else:
          for feature_idx in range(len(features)):
              data_feature = reduce_fn(data[feature_idx])
              mask[feature_idx].data = mask_fn(
                  data_feature, **sparse_config
              )
  
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 8 for calling the ActivationSparsifier obj's underlying member method code with example usage: obj.update_mask(name, data, configs) without return types