Adding a new Module for unsupported layer. Adding test for unsupported layers. Simple logging for unsupported layers (#1505)

Summary:

We are adding test for unsupported gradient layers. Open to ideas if there is a better way to structure the test.

A bit uncomfortable with removing pyre type validations as we allow anything to be passed into the GradientUnsupportedLayerOutput class.

Differential Revision: D69792994

Author: styusuf

captum/testing/helpers/basic_models.py

@@ -428,8 +428,27 @@ def forward(self, arg1: Tensor, arg2: Tensor) -> Tuple[Tensor, Tensor]:
         return (self.relu1(arg1), self.relu2(arg2))
 
 
+class GradientUnsupportedLayerOutput(nn.Module):
+    # pyre-fixme[2]: Parameter must be annotated.
+    def __init__(self, unsupported_layer_output) -> None:
+        super().__init__()
+        # pyre-fixme[4]: Attribute must be annotated.
+        self.unsupported_layer_output = unsupported_layer_output
+
+    @no_type_check
+    # pyre-fixme[3]: Return type must be annotated.
+    def forward(self):
+        return self.unsupported_layer_output
+
+
 class BasicModel_MultiLayer(nn.Module):
-    def __init__(self, inplace: bool = False, multi_input_module: bool = False) -> None:
+    def __init__(
+        self,
+        inplace: bool = False,
+        multi_input_module: bool = False,
+        # pyre-fixme[2]: Parameter must be annotated.
+        unsupported_layer_output=None,
+    ) -> None:
         super().__init__()
         # Linear 0 is simply identity transform
         self.multi_input_module = multi_input_module
@@ -445,6 +464,7 @@ def __init__(self, inplace: bool = False, multi_input_module: bool = False) -> N
         self.linear1_alt.bias = nn.Parameter(torch.tensor([-10.0, 1.0, 1.0, 1.0]))
         self.multi_relu = MultiRelu(inplace=inplace)
         self.relu = nn.ReLU(inplace=inplace)
+        self.unsupportedLayer = GradientUnsupportedLayerOutput(unsupported_layer_output)
 
         self.linear2 = nn.Linear(4, 2)
         self.linear2.weight = nn.Parameter(torch.ones(2, 4))
@@ -461,6 +481,9 @@ def forward(
         input = x if add_input is None else x + add_input
         lin0_out = self.linear0(input)
         lin1_out = self.linear1(lin0_out)
+        if self.unsupportedLayer is not None:
+            self.unsupportedLayer()
+            ## unsupportedLayer is unused intentionally for testing
         if self.multi_input_module:
             relu_out1, relu_out2 = self.multi_relu(lin1_out, self.linear1_alt(input))
             relu_out = relu_out1 + relu_out2