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:
      input_tensor (Tensor): Tensor to be gated with stochastic gates
  
  
  Returns:
      tuple[Tensor, Tensor]:
  
      - gated_input (Tensor): Tensor of the same shape weighted by the sampled
          gate values
  
      - l0_reg (Tensor): L0 regularization term to be optimized together with
          model loss,
          e.g. loss(model_out, target) + l0_reg
docstr_tok: ["Args",":","input_tensor","(","Tensor",")",":","Tensor","to","be","gated","with","stochastic","gates","Returns",":","tuple","[","Tensor",",","Tensor","]",":","-","gated_input","(","Tensor",")",":","Tensor","of","the","same","shape","weighted","by","the","sampled","gate","values","-","l0_reg","(","Tensor",")",":","L0","regularization","term","to","be","optimized","together","with","model","loss",",","e.g",".","loss","(","model_out",",","target",")","+","l0_reg"]
code: def forward(self, input_tensor: Tensor) -> Tuple[Tensor, Tensor]:
      """
      Args:
          input_tensor (Tensor): Tensor to be gated with stochastic gates
  
  
      Returns:
          tuple[Tensor, Tensor]:
  
          - gated_input (Tensor): Tensor of the same shape weighted by the sampled
              gate values
  
          - l0_reg (Tensor): L0 regularization term to be optimized together with
              model loss,
              e.g. loss(model_out, target) + l0_reg
      """
      if self.mask is None:
          n_ele = self._get_numel_of_input(input_tensor)
          assert n_ele == self.n_gates, (
              "if mask is not given, each example in the input batch should have the"
              " same number of elements"
              f" (received {n_ele}) as gates ({self.n_gates})"
          )
  
      input_size = input_tensor.size()
      batch_size = input_size[0]
  
      gate_values = self._sample_gate_values(batch_size)
  
      # hard-sigmoid rectification z=min(1,max(0,_z))
      gate_values = torch.clamp(gate_values, min=0, max=1)
  
      if self.mask is not None:
          # use expand_as not expand/broadcast_to which do not work with torch.fx
          input_mask = self.mask.expand_as(input_tensor)
  
          # flatten all dim except batch to gather from gate values
          flattened_mask = input_mask.reshape(batch_size, -1)
          gate_values = torch.gather(gate_values, 1, flattened_mask)
  
      # reshape gates(batch_size, n_elements) into input_size for point-wise mul
      gate_values = gate_values.reshape(input_size)
      gated_input = input_tensor * gate_values
  
      prob_density = self._get_gate_active_probs()
      if self.reg_reduction == "sum":
          l0_reg = prob_density.sum()
      elif self.reg_reduction == "mean":
          l0_reg = prob_density.mean()
      else:
          l0_reg = prob_density
  
      l0_reg *= self.reg_weight
  
      return gated_input, l0_reg
  
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 2 for calling the StochasticGatesBase obj's underlying member method code with example usage: obj.forward(input_tensor) and returns: gated_input,  l0_reg