Skip externalize for submodules not invoked in the exported graph (#18)

When a model registers a submodule whose class matches an
externalize_modules target but the model's forward never calls it, the
re-exported FX graph contains no custom-op node for that submodule.
The pipeline previously raised
``ValueError: Custom op for '<name>' not found in any ancestor program``
from _find_program_for, aborting the whole conversion.

Skip these submodules (with a UserWarning) instead. The unused module's
forward is still restored by _restore_externalized.

Also silence pre-existing F821 lint warnings on transform_with_custom_compression_ops
references in skipped tests so pre-commit ruff-check passes.

Co-authored-by: gokulkrishna98 <gokulkrishna98@users.noreply.github.com>

Author: gokulkrishna98

coreai_torch/_utils.py

@@ -1596,15 +1596,15 @@ def _find_program_for(
     name: str,
     op_name: str,
     programs: dict[str, ExportedProgram],
-) -> ExportedProgram:
-    """Find the nearest ancestor program that contains *op_name*."""
+) -> ExportedProgram | None:
+    """Find the nearest ancestor program that contains *op_name*, or ``None``."""
     for ancestor in _ancestor_paths(name):
         if (
             ancestor in programs
             and _find_custom_op_node(programs[ancestor], op_name) is not None
         ):
             return programs[ancestor]
-    raise ValueError(f"Custom op for '{name}' not found in any ancestor program")
+    return None
 
 
 def _reverse_lookup(mapping: dict[str, str], value: str) -> str | None:

coreai_torch/externalize.py

@@ -431,6 +431,18 @@ def __iter__(self) -> Iterator[_PreparedModule]:
             name: str = mod._externalize_name  # type: ignore[attr-defined]
             op_name: str = mod._externalize_op_name  # type: ignore[attr-defined]
             parent_ep = _find_program_for(name, op_name, self._programs)
+            if parent_ep is None:
+                # Marked submodule is not invoked in the exported graph.
+                # Skip it; its forward will be restored by _restore_externalized.
+                warnings.warn(
+                    f"\n[WARN] coreai_torch.externalize: skipping unused submodule '{name}'.\n"
+                    f"       It matched an externalize_modules target class but is not "
+                    f"reachable from the exported graph.\n"
+                    f"       Action: remove it from the model passed to add_pytorch_module, "
+                    f"or ignore if intentional.\n",
+                    stacklevel=2,
+                )
+                continue
             preps = _prepare_module_export(mod, parent_ep)
             for prep in preps:
                 prep._program_registry = self

tests/test_externalize.py

@@ -3683,3 +3683,108 @@ def forward(self, x: torch.Tensor, indices: torch.Tensor) -> torch.Tensor:
     await _validate_numerics(
         coreai_program, model, sample, input_names=("x", "indices")
     )
+
+
+@pytest.mark.ir
+def test_externalize_unused_submodule_ir() -> None:
+    """An externalizable submodule that the model's forward never calls is skipped.
+
+    Previously, a registered submodule that did not appear in the exported graph
+    caused the externalize pipeline to raise ``ValueError: Custom op for ...
+    not found in any ancestor program``. The pipeline should instead warn and
+    proceed, lowering the rest of the model normally.
+    """
+
+    class InnerModule(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.fc = nn.Linear(4, 4)
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            return torch.relu(self.fc(x))
+
+    class OuterModel(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.pre = nn.Linear(4, 4)
+            # Registered as a submodule but intentionally never invoked.
+            self.unused = InnerModule()
+            self.post = nn.Linear(4, 4)
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            return self.post(self.pre(x))
+
+    torch.manual_seed(42)
+    model = OuterModel().eval()
+    sample = (torch.randn(2, 4),)
+
+    with pytest.warns(UserWarning, match="skipping unused submodule"):
+        converter = TorchConverter().add_pytorch_module(
+            model,
+            export_fn=lambda m: torch.export.export(m, args=sample).run_decompositions(
+                get_decomp_table()
+            ),
+            externalize_modules=[InnerModule],
+        )
+        coreai_program = converter.to_coreai()
+
+    check_file = """
+        // CHECK-LABEL: module {
+        // CHECK-NOT: coreai.graph noinline
+        // CHECK: coreai.graph @main(
+        // CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<2x4xf32> {coreai.name = "x"}
+        // CHECK-SAME: ) -> (tensor<2x4xf32>
+        // CHECK: %[[MM0:[0-9a-z_]+]] = coreai.decomposable.broadcasting_batch_matmul %[[ARG0]],
+        // CHECK-SAME: : (tensor<2x4xf32>, tensor<4x4xf32>) -> tensor<2x4xf32>
+        // CHECK: %[[ADD0:[0-9a-z_]+]] = coreai.decomposable.broadcasting_add %[[MM0]],
+        // CHECK-SAME: : (tensor<2x4xf32>, tensor<4xf32>) -> tensor<2x4xf32>
+        // CHECK: %[[MM1:[0-9a-z_]+]] = coreai.decomposable.broadcasting_batch_matmul %[[ADD0]],
+        // CHECK-SAME: : (tensor<2x4xf32>, tensor<4x4xf32>) -> tensor<2x4xf32>
+        // CHECK: %[[ADD1:[0-9a-z_]+]] = coreai.decomposable.broadcasting_add %[[MM1]],
+        // CHECK-SAME: : (tensor<2x4xf32>, tensor<4xf32>) -> tensor<2x4xf32>
+        // CHECK-NOT: coreai.invoke
+        // CHECK-NOT: coreai.relu
+        // CHECK: coreai.output %[[ADD1]] : tensor<2x4xf32>
+        // CHECK: }
+        // CHECK-NOT: coreai.graph
+        // CHECK: }
+    """
+    filecheck_pattern(str(coreai_program), check_file=check_file)
+
+
+async def test_externalize_unused_submodule_numerics() -> None:
+    """Numerics: unused externalizable submodule does not affect output."""
+
+    class InnerModule(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.fc = nn.Linear(4, 4)
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            return torch.relu(self.fc(x))
+
+    class OuterModel(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.pre = nn.Linear(4, 4)
+            self.unused = InnerModule()
+            self.post = nn.Linear(4, 4)
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            return self.post(self.pre(x))
+
+    torch.manual_seed(42)
+    model = OuterModel().eval()
+    sample = (torch.randn(2, 4),)
+
+    with pytest.warns(UserWarning, match="skipping unused submodule"):
+        converter = TorchConverter().add_pytorch_module(
+            model,
+            export_fn=lambda m: torch.export.export(m, args=sample).run_decompositions(
+                get_decomp_table()
+            ),
+            externalize_modules=[InnerModule],
+        )
+        coreai_program = converter.to_coreai()
+
+    await _validate_numerics(coreai_program, model, sample)