Add pyre fixme for downstream type errors from D65753120 (#1439)

Summary:
Pull Request resolved: #1439

X-link: facebook/Ax#3055

X-link: ctrl-labs/src2#38515

X-link: ctrl-labs/src2#38514

bypass-github-export-checks "The check is bugged, I exported all the required exports"

Reviewed By: jermenkoo

Differential Revision: D65826205

fbshipit-source-id: b08f00caaac8f3cc235bc280c4fe3f99089a0753

Author: ezyang

captum/attr/_core/feature_ablation.py

@@ -806,9 +806,11 @@ def _construct_ablated_input(
             dim=0,
         ).long()
         current_mask = current_mask.to(expanded_input.device)
-        assert baseline is not None, "baseline must be provided"
         ablated_tensor = (
-            expanded_input * (1 - current_mask).to(expanded_input.dtype)
+            expanded_input
+            * (1 - current_mask).to(expanded_input.dtype)
+            # pyre-fixme[58]: `*` is not supported for operand types `Union[None, float,
+            #  Tensor]` and `Tensor`.
         ) + (baseline * current_mask.to(expanded_input.dtype))
         return ablated_tensor, current_mask
 

captum/attr/_core/occlusion.py

@@ -316,13 +316,14 @@ def _construct_ablated_input(
             ],
             dim=0,
         ).long()
-        assert baseline is not None, "baseline should not be None"
         ablated_tensor = (
             expanded_input
             * (
                 torch.ones(1, dtype=torch.long, device=expanded_input.device)
                 - input_mask
             ).to(expanded_input.dtype)
+            # pyre-fixme[58]: `*` is not supported for operand types `Union[None, float,
+            #  Tensor]` and `Tensor`.
         ) + (baseline * input_mask.to(expanded_input.dtype))
         return ablated_tensor, input_mask
 

captum/module/gaussian_stochastic_gates.py

@@ -133,9 +133,9 @@ def _get_gate_active_probs(self) -> Tensor:
             probs (Tensor): probabilities tensor of the gates are active
                 in shape(n_gates)
         """
-        std = self.std
-        assert std is not None, "std should not be None"
-        x = self.mu / std
+        # pyre-fixme[58]: `/` is not supported for operand types `Parameter` and
+        #  `Optional[float]`.
+        x = self.mu / self.std
         return 0.5 * (1 + torch.erf(x / math.sqrt(2)))
 
     @classmethod

tests/attr/test_input_x_gradient.py

@@ -117,9 +117,9 @@ def _input_x_gradient_classification_assert(self, nt_type: str = "vanilla") -> N
             attributions = input_x_grad.attribute(input, target)
             output = model(input)[:, target]
             output.backward()
-            input_grad = input.grad
-            assert input_grad is not None
-            expected = input_grad * input
+            # pyre-fixme[58]: `*` is not supported for operand types
+            #  `Optional[Tensor]` and `Tensor`.
+            expected = input.grad * input
             assertTensorAlmostEqual(self, attributions, expected, 0.00001, "max")
         else:
             nt = NoiseTunnel(input_x_grad)