class_docstr: Abstract base class for creation of new pruning techniques.
  
  Provides a skeleton for customization requiring the overriding of methods
  such as :meth:`compute_mask` and :meth:`apply`.
class_docstr_tok: ["Abstract","base","class","for","creation","of","new","pruning","techniques",".","Provides","a","skeleton","for","customization","requiring","the","overriding","of","methods","such","as",":","meth",":","`","compute_mask","`","and",":","meth",":","`","apply","`","."]
docstr: Add pruning on the fly and reparametrization of a tensor.
  
  Adds the forward pre-hook that enables pruning on the fly and
  the reparametrization of a tensor in terms of the original tensor
  and the pruning mask.
  
  Args:
      module (nn.Module): module containing the tensor to prune
      name (str): parameter name within ``module`` on which pruning
          will act.
      args: arguments passed on to a subclass of
          :class:`BasePruningMethod`
      importance_scores (torch.Tensor): tensor of importance scores (of
          same shape as module parameter) used to compute mask for pruning.
          The values in this tensor indicate the importance of the
          corresponding elements in the parameter being pruned.
          If unspecified or None, the parameter will be used in its place.
      kwargs: keyword arguments passed on to a subclass of a
          :class:`BasePruningMethod`
docstr_tok: ["Add","pruning","on","the","fly","and","reparametrization","of","a","tensor",".","Adds","the","forward","pre-hook","that","enables","pruning","on","the","fly","and","the","reparametrization","of","a","tensor","in","terms","of","the","original","tensor","and","the","pruning","mask",".","Args",":","module","(","nn.Module",")",":","module","containing","the","tensor","to","prune","name","(","str",")",":","parameter","name","within","``","module","``","on","which","pruning","will","act",".","args",":","arguments","passed","on","to","a","subclass","of",":","class",":","`","BasePruningMethod","`","importance_scores","(","torch.Tensor",")",":","tensor","of","importance","scores","(","of","same","shape","as","module","parameter",")","used","to","compute","mask","for","pruning",".","The","values","in","this","tensor","indicate","the","importance","of","the","corresponding","elements","in","the","parameter","being","pruned",".","If","unspecified","or","None",",","the","parameter","will","be","used","in","its","place",".","kwargs",":","keyword","arguments","passed","on","to","a","subclass","of","a",":","class",":","`","BasePruningMethod","`"]
code: def apply(cls, module, name, *args, importance_scores=None, **kwargs):
      r"""Add pruning on the fly and reparametrization of a tensor.
  
      Adds the forward pre-hook that enables pruning on the fly and
      the reparametrization of a tensor in terms of the original tensor
      and the pruning mask.
  
      Args:
          module (nn.Module): module containing the tensor to prune
          name (str): parameter name within ``module`` on which pruning
              will act.
          args: arguments passed on to a subclass of
              :class:`BasePruningMethod`
          importance_scores (torch.Tensor): tensor of importance scores (of
              same shape as module parameter) used to compute mask for pruning.
              The values in this tensor indicate the importance of the
              corresponding elements in the parameter being pruned.
              If unspecified or None, the parameter will be used in its place.
          kwargs: keyword arguments passed on to a subclass of a
              :class:`BasePruningMethod`
      """
  
      def _get_composite_method(cls, module, name, *args, **kwargs):
          # Check if a pruning method has already been applied to
          # `module[name]`. If so, store that in `old_method`.
          old_method = None
          found = 0
          # there should technically be only 1 hook with hook.name == name
          # assert this using `found`
          hooks_to_remove = []
          for k, hook in module._forward_pre_hooks.items():
              # if it exists, take existing thing, remove hook, then
              # go through normal thing
              if (
                  isinstance(hook, BasePruningMethod)
                  and hook._tensor_name == name
              ):
                  old_method = hook
                  hooks_to_remove.append(k)
                  found += 1
          assert (
              found <= 1
          ), f"Avoid adding multiple pruning hooks to the\
                  same tensor {name} of module {module}. Use a PruningContainer."
  
          for k in hooks_to_remove:
              del module._forward_pre_hooks[k]
  
          # Apply the new pruning method, either from scratch or on top of
          # the previous one.
          method = cls(*args, **kwargs)  # new pruning
          # Have the pruning method remember what tensor it's been applied to
          method._tensor_name = name
  
          # combine `methods` with `old_method`, if `old_method` exists
          if old_method is not None:  # meaning that there was a hook
              # if the hook is already a pruning container, just add the
              # new pruning method to the container
              if isinstance(old_method, PruningContainer):
                  old_method.add_pruning_method(method)
                  method = old_method  # rename old_method --> method
  
              # if the hook is simply a single pruning method, create a
              # container, add the old pruning method and the new one
              elif isinstance(old_method, BasePruningMethod):
                  container = PruningContainer(old_method)
                  # Have the pruning method remember the name of its tensor
                  # setattr(container, '_tensor_name', name)
                  container.add_pruning_method(method)
                  method = container  # rename container --> method
          return method
  
      method = _get_composite_method(cls, module, name, *args, **kwargs)
      # at this point we have no forward_pre_hooks but we could have an
      # active reparametrization of the tensor if another pruning method
      # had been applied (in which case `method` would be a PruningContainer
      # and not a simple pruning method).
  
      # Pruning is to be applied to the module's tensor named `name`,
      # starting from the state it is found in prior to this iteration of
      # pruning. The pruning mask is calculated based on importances scores.
  
      orig = getattr(module, name)
      if importance_scores is not None:
          assert (
              importance_scores.shape == orig.shape
          ), f"importance_scores should have the same shape as parameter                 {name} of {module}"
      else:
          importance_scores = orig
  
      # If this is the first time pruning is applied, take care of moving
      # the original tensor to a new parameter called name + '_orig' and
      # and deleting the original parameter
      if not isinstance(method, PruningContainer):
          # copy `module[name]` to `module[name + '_orig']`
          module.register_parameter(name + "_orig", orig)
          # temporarily delete `module[name]`
          del module._parameters[name]
          default_mask = torch.ones_like(orig)  # temp
      # If this is not the first time pruning is applied, all of the above
      # has been done before in a previous pruning iteration, so we're good
      # to go
      else:
          default_mask = (
              getattr(module, name + "_mask")
              .detach()
              .clone(memory_format=torch.contiguous_format)
          )
  
      # Use try/except because if anything goes wrong with the mask
      # computation etc., you'd want to roll back.
      try:
          # get the final mask, computed according to the specific method
          mask = method.compute_mask(
              importance_scores, default_mask=default_mask
          )
          # reparameterize by saving mask to `module[name + '_mask']`...
          module.register_buffer(name + "_mask", mask)
          # ... and the new pruned tensor to `module[name]`
          setattr(module, name, method.apply_mask(module))
          # associate the pruning method to the module via a hook to
          # compute the function before every forward() (compile by run)
          module.register_forward_pre_hook(method)
  
      except Exception as e:
          if not isinstance(method, PruningContainer):
              orig = getattr(module, name + "_orig")
              module.register_parameter(name, orig)
              del module._parameters[name + "_orig"]
          raise e
  
      return method
  
imports: import numbers
  from abc import ABC, abstractmethod
  from collections.abc import Iterable
  from typing import Tuple
  import torch
label_desc: Use image node_id 4 for calling the BasePruningMethod obj's underlying member method code with example usage: obj.apply(cls, module, name) and returns: method, method