Adding a new Module for unsupported layer. Adding test for unsupported layers. Simple logging for unsupported layers

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

Differential Revision: D69792994

Author: styusuf

captum/testing/helpers/basic_models.py

@@ -2,7 +2,7 @@
 
 # pyre-strict
 
-from typing import no_type_check, Optional, Tuple, Union
+from typing import List, no_type_check, Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
@@ -428,8 +428,24 @@ def forward(self, arg1: Tensor, arg2: Tensor) -> Tuple[Tensor, Tensor]:
         return (self.relu1(arg1), self.relu2(arg2))
 
 
+class GradientUnsupportedLayerOutput(nn.Module):
+    def __init__(self, inplace: bool = False) -> None:
+        super().__init__()
+        self.relu1 = nn.ReLU(inplace=inplace)
+        self.relu2 = nn.ReLU(inplace=inplace)
+
+    @no_type_check
+    def forward(self, arg1: Tensor, arg2: Tensor) -> List[Tensor]:
+        return [self.relu1(arg1), self.relu2(arg2)]
+
+
 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,
+        unsupported_layer_output: bool = False,
+    ) -> None:
         super().__init__()
         # Linear 0 is simply identity transform
         self.multi_input_module = multi_input_module
@@ -443,7 +459,11 @@ def __init__(self, inplace: bool = False, multi_input_module: bool = False) -> N
         self.linear1_alt = nn.Linear(3, 4)
         self.linear1_alt.weight = nn.Parameter(torch.ones(4, 3))
         self.linear1_alt.bias = nn.Parameter(torch.tensor([-10.0, 1.0, 1.0, 1.0]))
-        self.multi_relu = MultiRelu(inplace=inplace)
+        self.multi_relu = (
+            GradientUnsupportedLayerOutput(inplace=inplace)
+            if unsupported_layer_output
+            else MultiRelu(inplace=inplace)
+        )
         self.relu = nn.ReLU(inplace=inplace)
 
         self.linear2 = nn.Linear(4, 2)