class_docstr: Abstract module for Stochastic Gates.
  
  Stochastic Gates is a practical solution to add L0 norm regularization for neural
  networks. L0 regularization, which explicitly penalizes any present (non-zero)
  parameters, can help network pruning and feature selection, but directly optimizing
  L0 is a non-differentiable combinatorial problem. To surrogate L0, Stochastic Gate
  uses certain continuous probability distributions (e.g., Concrete, Gaussian) with
  hard-sigmoid rectification as a continuous smoothed Bernoulli distribution
  determining the weight of a parameter, i.e., gate. Then L0 is equal to the gates's
  non-zero probability represented by the parameters of the continuous probability
  distribution. The gate value can also be reparameterized to the distribution
  parameters with a noise. So the expected L0 can be optimized through learning
  the distribution parameters via stochastic gradients.
  
  This base class defines the shared variables and forward logic of how the input is
  gated regardless of the underneath distribution. The actual implementation should
  extend this class and implement the distribution specific functions.
docstr: Args:
      n_gates (int): number of gates.
  
      mask (Tensor, optional): If provided, this allows grouping multiple
          input tensor elements to share the same stochastic gate.
          This tensor should be broadcastable to match the input shape
          and contain integers in the range 0 to n_gates - 1.
          Indices grouped to the same stochastic gate should have the same value.
          If not provided, each element in the input tensor
          (on dimensions other than dim 0 - batch dim) is gated separately.
          Default: None
  
      reg_weight (float, optional): rescaling weight for L0 regularization term.
          Default: 1.0
  
      reg_reduction (str, optional): the reduction to apply to the regularization:
          ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be
          applied and it will be the same as the return of ``get_active_probs``,
          ``'mean'``: the sum of the gates non-zero probabilities will be divided
          by the number of gates, ``'sum'``: the gates non-zero probabilities will
          be summed.
          Default: ``'sum'``
docstr_tok: ["Args",":","n_gates","(","int",")",":","number","of","gates",".","mask","(","Tensor",",","optional",")",":","If","provided",",","this","allows","grouping","multiple","input","tensor","elements","to","share","the","same","stochastic","gate",".","This","tensor","should","be","broadcastable","to","match","the","input","shape","and","contain","integers","in","the","range","0","to","n_gates","-","1",".","Indices","grouped","to","the","same","stochastic","gate","should","have","the","same","value",".","If","not","provided",",","each","element","in","the","input","tensor","(","on","dimensions","other","than","dim","0","-","batch","dim",")","is","gated","separately",".","Default",":","None","reg_weight","(","float",",","optional",")",":","rescaling","weight","for","L0","regularization","term",".","Default",":","1.0","reg_reduction","(","str",",","optional",")",":","the","reduction","to","apply","to","the","regularization",":","``","'none","'","``","|","``","'mean","'","``","|","``","'sum","'","``",".","``","'none","'","``",":","no","reduction","will","be","applied","and","it","will","be","the","same","as","the","return","of","``","get_active_probs","``",",","``","'mean","'","``",":","the","sum","of","the","gates","non-zero","probabilities","will","be","divided","by","the","number","of","gates",",","``","'sum","'","``",":","the","gates","non-zero","probabilities","will","be","summed",".","Default",":","``","'sum","'","``"]
code: def __init__(
      self,
      n_gates: int,
      mask: Optional[Tensor] = None,
      reg_weight: float = 1.0,
      reg_reduction: str = "sum",
  ):
      """
      Args:
          n_gates (int): number of gates.
  
          mask (Tensor, optional): If provided, this allows grouping multiple
              input tensor elements to share the same stochastic gate.
              This tensor should be broadcastable to match the input shape
              and contain integers in the range 0 to n_gates - 1.
              Indices grouped to the same stochastic gate should have the same value.
              If not provided, each element in the input tensor
              (on dimensions other than dim 0 - batch dim) is gated separately.
              Default: None
  
          reg_weight (float, optional): rescaling weight for L0 regularization term.
              Default: 1.0
  
          reg_reduction (str, optional): the reduction to apply to the regularization:
              ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be
              applied and it will be the same as the return of ``get_active_probs``,
              ``'mean'``: the sum of the gates non-zero probabilities will be divided
              by the number of gates, ``'sum'``: the gates non-zero probabilities will
              be summed.
              Default: ``'sum'``
      """
      super().__init__()
  
      if mask is not None:
          max_mask_ind = mask.max().item()
          assert max_mask_ind == n_gates - 1, (
              f"the maximum mask index (received {max_mask_ind}) should be equal to"
              f" the number of gates - 1 (received {n_gates}) since each mask"
              " should correspond to a gate"
          )
  
      valid_reg_reduction = ["none", "mean", "sum"]
      assert (
          reg_reduction in valid_reg_reduction
      ), f"reg_reduction must be one of [none, mean, sum], received: {reg_reduction}"
      self.reg_reduction = reg_reduction
  
      self.n_gates = n_gates
      self.register_buffer(
          "mask", mask.detach().clone() if mask is not None else None
      )
      self.reg_weight = reg_weight
  
file_path: captum/captum/module/stochastic_gates_base.py
imports: from abc import ABC, abstractmethod
  from typing import Optional, Tuple
  import torch
  from torch import Tensor
  from torch.nn import Module
label_desc: Use image node_id 1 to create a new StochasticGatesBase object from inherited base classes: Module, ABC with example: obj = StochasticGatesBase(n_gates, mask, reg_weight, reg_reduction)