Adding test for output that is a tensor of integers. Updating passthrough layer. (#1526)

Summary:
Pull Request resolved: #1526

We are adding tests for different types of unsupported and non-differentiable layer output. Here, we add a test for layer output that is a tensor of integers.

We split by the cases for unsupported layers from the case when the layer output is used by some tasks and not others.

When layer output is not supported (layer output is a List of Tensors or a Tensor of integers), we don't get attributions and return None for those layers.

In the case when a layer output is not used by a task, we should output a tensor of zeros for that task.

Reviewed By: craymichael

Differential Revision: D70919347

fbshipit-source-id: 191d9d69c78bcf00fa3cbbbd5707154e0f221410

Author: styusuf

captum/testing/helpers/basic_models.py

@@ -418,7 +418,7 @@ def forward(self, input1, input2, input3=None):
         return self.linear2(self.relu(self.linear1(embeddings))).sum(1)
 
 
-class GradientUnsupportedLayerOutput(nn.Module):
+class PassThroughLayerOutput(nn.Module):
     """
     This layer is used to test the case where the model returns a layer that
     is not supported by the gradient computation.
@@ -428,10 +428,8 @@ def __init__(self) -> None:
         super().__init__()
 
     @no_type_check
-    def forward(
-        self, unsupported_layer_output: PassThroughOutputType
-    ) -> PassThroughOutputType:
-        return unsupported_layer_output
+    def forward(self, output: PassThroughOutputType) -> PassThroughOutputType:
+        return output
 
 
 class BasicModel_GradientLayerAttribution(nn.Module):
@@ -456,7 +454,7 @@ def __init__(
 
         self.relu = nn.ReLU(inplace=inplace)
         self.relu_alt = nn.ReLU(inplace=False)
-        self.unsupportedLayer = GradientUnsupportedLayerOutput()
+        self.unsupported_layer = PassThroughLayerOutput()
 
         self.linear2 = nn.Linear(4, 2)
         self.linear2.weight = nn.Parameter(torch.ones(2, 4))
@@ -466,6 +464,8 @@ def __init__(
         self.linear3.weight = nn.Parameter(torch.ones(2, 4))
         self.linear3.bias = nn.Parameter(torch.tensor([-1.0, 1.0]))
 
+        self.int_layer = PassThroughLayerOutput()  # sample layer with an int ouput
+
     @no_type_check
     def forward(
         self, x: Tensor, add_input: Optional[Tensor] = None
@@ -476,7 +476,7 @@ def forward(
         lin1_out_alt = self.linear1_alt(lin0_out)
 
         if self.unsupported_layer_output is not None:
-            self.unsupportedLayer(self.unsupported_layer_output)
+            self.unsupported_layer(self.unsupported_layer_output)
             # unsupportedLayer is unused in the forward func.
         self.relu_alt(
             lin1_out_alt
@@ -485,9 +485,10 @@ def forward(
         relu_out = self.relu(lin1_out)
         lin2_out = self.linear2(relu_out)
 
-        lin3_out = self.linear3(lin1_out_alt).to(torch.int64)
+        lin3_out = self.linear3(lin1_out_alt)
+        int_output = self.int_layer(lin3_out.to(torch.int64))
 
-        output_tensors = torch.cat((lin2_out, lin3_out), dim=1)
+        output_tensors = torch.cat((lin2_out, int_output), dim=1)
 
         # we return a dictionary of tensors as an output to test the case
         # where an output accessor is required