Fix is_layer_at_idx for LRP (#308)

* Fix `is_layer_at_idx` for LRP
* add composite lrp test

src/innvestigate/analyzer/relevance_based/relevance_analyzer.py

@@ -394,7 +394,9 @@ def __init__(
         # apply rule to first self._until_layer_idx layers
         if self._until_layer_rule is not None and self._until_layer_idx is not None:
             for i in range(self._until_layer_idx + 1):
-                is_at_idx: LayerCheck = lambda layer: ichecks.is_layer_at_idx(layer, i)
+                is_at_idx: LayerCheck = lambda layer, i=i: ichecks.is_layer_at_idx(
+                    model, layer, i
+                )
                 rules.insert(0, (is_at_idx, self._until_layer_rule))
 
         # create a BoundedRule for input layer handling from given tuple

src/innvestigate/backend/checks.py

@@ -7,7 +7,7 @@
 from tensorflow import Module, keras
 
 import innvestigate.backend as ibackend
-from innvestigate.backend.types import Layer
+from innvestigate.backend.types import Layer, Model
 
 __all__ = [
     "get_activation_search_safe_layers",
@@ -325,8 +325,6 @@ def is_input_layer(layer: Layer, ignore_reshape_layers: bool = True) -> bool:
     return all(isinstance(x, klayers.InputLayer) for x in layer_inputs)
 
 
-def is_layer_at_idx(layer: Layer, index, ignore_reshape_layers=True) -> bool:
-    """Checks if layer is a layer at index index,
-    by repeatedly applying is_input_layer()."""
-    # TODO: implement layer index check
-    raise NotImplementedError("Layer index checking hasn't been implemented yet.")
+def is_layer_at_idx(model: Model, layer: Layer, index) -> bool:
+    """Checks if layer is a layer at specified index of model."""
+    return layer == model.layers[index]

tests/backend/test_layer_idx.py

@@ -0,0 +1,42 @@
+import pytest
+import tensorflow.keras.layers as klayers
+import tensorflow.keras.models as kmodels
+
+from innvestigate.analyzer.relevance_based.relevance_analyzer import LRP
+
+
+@pytest.mark.graph
+@pytest.mark.fast
+@pytest.mark.lrp
+def test_composite_lrp():
+    model = kmodels.Sequential(
+        [
+            klayers.Input(shape=(28, 28, 3)),
+            klayers.Conv2D(8, 3, activation="relu"),
+            klayers.Conv2D(4, 5, activation="relu"),
+            klayers.Flatten(),
+            klayers.Dense(16, activation="relu"),
+            klayers.Dense(2, activation="softmax"),
+        ]
+    )
+    analyzer = LRP(
+        model,
+        rule="Z",
+        input_layer_rule="Flat",
+        until_layer_idx=2,
+        until_layer_rule="Epsilon",
+    )
+    correct_rules = [
+        "Flat",
+        "Epsilon",
+        "Epsilon",
+        "Z",
+        "Z",
+    ]  # Correct rules corresponding to analyzer input args
+
+    for i, layer in enumerate(model.layers):
+        for condition, rule in analyzer._rules:
+            if condition(layer):
+                rule_class = rule
+                break
+        assert rule_class == correct_rules[i]