Detect nested Keras Lambda layers

modelscan/scanners/h5/scan.py

@@ -19,6 +19,7 @@
 from modelscan.scanners.saved_model.scan import SavedModelLambdaDetectScan
 from modelscan.model import Model
 from modelscan.settings import SupportedModelFormats
+from modelscan.scanners.keras_utils import get_keras_layer_names
 
 logger = logging.getLogger("modelscan")
 
@@ -111,23 +112,18 @@ def _get_keras_h5_operator_names(self, model: Model) -> Optional[List[Any]]:
                     return None
 
                 model_config = json.loads(model_hdf5.attrs.get("model_config", {}))
-                layers = model_config.get("config", {}).get("layers", {})
-                lambda_layers = []
-                for layer in layers:
-                    if layer.get("class_name", {}) == "Lambda":
-                        lambda_layers.append(
-                            layer.get("config", {}).get("function", {})
-                        )
+                lambda_layers = [
+                    layer_name
+                    for layer_name in get_keras_layer_names(model_config)
+                    if layer_name == "Lambda"
+                ]
             except json.JSONDecodeError as e:
                 logger.error(
                     f"Not a valid JSON data from source: {model.get_source()}, error: {e}"
                 )
                 return ["JSONDecodeError"]
 
-        if lambda_layers:
-            return ["Lambda"] * len(lambda_layers)
-
-        return []
+        return lambda_layers
 
     def handle_binary_dependencies(
         self, settings: Optional[Dict[str, Any]] = None

modelscan/scanners/keras/scan.py

@@ -10,6 +10,7 @@
 from modelscan.scanners.saved_model.scan import SavedModelLambdaDetectScan
 from modelscan.model import Model
 from modelscan.settings import SupportedModelFormats
+from modelscan.scanners.keras_utils import get_keras_layer_names
 
 
 logger = logging.getLogger("modelscan")
@@ -119,15 +120,11 @@ def _scan_keras_config_file(self, model: Model) -> ScanResults:
     def _get_keras_operator_names(self, model: Model) -> List[str]:
         model_config_data = json.load(model.get_stream())
 
-        lambda_layers = [
-            layer.get("config", {}).get("function", {})
-            for layer in model_config_data.get("config", {}).get("layers", {})
-            if layer.get("class_name", {}) == "Lambda"
+        return [
+            layer_name
+            for layer_name in get_keras_layer_names(model_config_data)
+            if layer_name == "Lambda"
         ]
-        if lambda_layers:
-            return ["Lambda"] * len(lambda_layers)
-
-        return []
 
     @staticmethod
     def name() -> str:

modelscan/scanners/keras_utils.py

@@ -0,0 +1,19 @@
+from typing import Any, List
+
+
+def get_keras_layer_names(config: Any) -> List[str]:
+    layer_names: List[str] = []
+
+    if isinstance(config, dict):
+        class_name = config.get("class_name")
+        if isinstance(class_name, str):
+            layer_names.append(class_name)
+
+        for value in config.values():
+            layer_names.extend(get_keras_layer_names(value))
+
+    elif isinstance(config, list):
+        for item in config:
+            layer_names.extend(get_keras_layer_names(item))
+
+    return layer_names

tests/test_keras_nested_lambda_scan.py

@@ -0,0 +1,84 @@
+import io
+import json
+import zipfile
+from types import SimpleNamespace
+
+from modelscan.model import Model
+from modelscan.issues import IssueSeverity
+from modelscan.scanners.keras.scan import KerasLambdaDetectScan
+from modelscan.settings import DEFAULT_SETTINGS, SupportedModelFormats
+
+
+def test_keras_scan_detects_nested_lambda_layers() -> None:
+    model_config = {
+        "class_name": "Functional",
+        "config": {
+            "layers": [
+                {
+                    "class_name": "Functional",
+                    "config": {
+                        "layers": [
+                            {
+                                "class_name": "InputLayer",
+                                "config": {"name": "nested_input"},
+                            },
+                            {
+                                "class_name": "Lambda",
+                                "config": {"name": "nested_lambda"},
+                            },
+                        ]
+                    },
+                }
+            ]
+        },
+    }
+    model = Model(
+        "nested.keras:config.json", io.BytesIO(json.dumps(model_config).encode())
+    )
+
+    assert KerasLambdaDetectScan({})._get_keras_operator_names(model) == ["Lambda"]
+
+
+def test_keras_scanner_reports_nested_lambda_layers(monkeypatch) -> None:  # type: ignore[no-untyped-def]
+    from modelscan.scanners.saved_model import scan as saved_model_scan
+
+    monkeypatch.setattr(saved_model_scan, "tensorflow_installed", True)
+    monkeypatch.setattr(
+        saved_model_scan,
+        "tensorflow",
+        SimpleNamespace(raw_ops=SimpleNamespace()),
+        raising=False,
+    )
+
+    model_config = {
+        "class_name": "Functional",
+        "config": {
+            "layers": [
+                {
+                    "class_name": "Functional",
+                    "config": {
+                        "layers": [
+                            {
+                                "class_name": "Lambda",
+                                "config": {"name": "nested_lambda"},
+                            }
+                        ]
+                    },
+                }
+            ]
+        },
+    }
+    archive = io.BytesIO()
+    with zipfile.ZipFile(archive, "w") as model_zip:
+        model_zip.writestr("config.json", json.dumps(model_config))
+    archive.seek(0)
+
+    model = Model("nested.keras", archive)
+    model.set_context("formats", [SupportedModelFormats.KERAS])
+
+    result = KerasLambdaDetectScan(DEFAULT_SETTINGS).scan(model)
+
+    assert result is not None
+    assert len(result.issues) == 1
+    assert result.issues[0].details.operator == "Lambda"
+    assert result.issues[0].severity == IssueSeverity.MEDIUM