update stg sphinx to include inherited methods

captum/module/binary_concrete_stochastic_gates.py

@@ -133,22 +133,6 @@ def __init__(
         # pre-calculate the fixed term used in active prob
         self.active_prob_offset = temperature * math.log(-lower_bound / upper_bound)
 
-    def forward(self, *args, **kwargs):
-        """
-        Args:
-            input_tensor (Tensor): Tensor to be gated with stochastic gates
-
-
-        Outputs:
-            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
-        """
-        return super().forward(*args, **kwargs)
-
     def _sample_gate_values(self, batch_size: int) -> Tensor:
         """
         Sample gate values for each example in the batch from the binary concrete

captum/module/gaussian_stochastic_gates.py

@@ -78,21 +78,6 @@ def __init__(
         assert 0 < std, f"the standard deviation should be positive, received {std}"
         self.std = std
 
-    def forward(self, *args, **kwargs):
-        """
-        Args:
-            input_tensor (Tensor): Tensor to be gated with stochastic gates
-
-        Outputs:
-            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
-        """
-        return super().forward(*args, **kwargs)
-
     def _sample_gate_values(self, batch_size: int) -> Tensor:
         """
         Sample gate values for each example in the batch from the Gaussian distribution

captum/module/stochastic_gates_base.py

@@ -87,11 +87,13 @@ def forward(self, input_tensor: Tensor) -> Tuple[Tensor, Tensor]:
             input_tensor (Tensor): Tensor to be gated with stochastic gates
 
 
-        Outputs:
-            gated_input (Tensor): Tensor of the same shape weighted by the sampled
+        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
+            - l0_reg (Tensor): L0 regularization term to be optimized together with
                 model loss,
                 e.g. loss(model_out, target) + l0_reg
         """
@@ -140,16 +142,18 @@ def get_gate_values(self, clamp: bool = True) -> Tensor:
         Get the gate values, which are the means of the underneath gate distributions,
         optionally clamped within 0 and 1.
 
-        Returns:
-            gate_values (Tensor): value of each gate in shape(n_gates)
-
-            clamp (bool): if clamp the gate values. As smoothed Bernoulli
-                variables, gate values are clamped withn 0 and 1 by defautl.
+        Args:
+            clamp (bool): whether to clamp the gate values or not. As smoothed Bernoulli
+                variables, gate values are clamped within 0 and 1 by default.
                 Turn this off to get the raw means of the underneath
-                distribution (e.g., conrete, gaussian), which can be useful to
+                distribution (e.g., concrete, gaussian), which can be useful to
                 differentiate the gates' importance when multiple gate
                 values are beyond 0 or 1.
                 Default: True
+
+        Returns:
+            Tensor:
+            - gate_values (Tensor): value of each gate in shape(n_gates)
         """
         gate_values = self._get_gate_values()
         if clamp:
@@ -162,7 +166,8 @@ def get_gate_active_probs(self) -> Tensor:
         Get the active probability of each gate, i.e, gate value > 0
 
         Returns:
-            probs (Tensor): probabilities tensor of the gates are active
+            Tensor:
+            - probs (Tensor): probabilities tensor of the gates are active
                 in shape(n_gates)
         """
         return self._get_gate_active_probs().detach()

sphinx/source/binary_concrete_stg.rst

@@ -3,3 +3,4 @@ BinaryConcreteStochasticGates
 
 .. autoclass:: captum.module.BinaryConcreteStochasticGates
     :members:
+    :inherited-members: Module

sphinx/source/gaussian_stg.rst

@@ -3,3 +3,4 @@ GaussianStochasticGates
 
 .. autoclass:: captum.module.GaussianStochasticGates
     :members:
+    :inherited-members: Module