Skip externalize for submodules not invoked in the exported graph

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.

Author: gokulkrishna98

coreai_torch/_utils.py

@@ -1575,15 +1575,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

@@ -3032,10 +3032,12 @@ def forward(
 
 @pytest.mark.ir
 @pytest.mark.flaky(reruns=3)
-@pytest.mark.skip(reason=(
-    "transform_with_custom_compression_ops has been deprecated. Consider removing "
-    "these tests or use an alternative way to generate quantized weights"
-))
+@pytest.mark.skip(
+    reason=(
+        "transform_with_custom_compression_ops has been deprecated. Consider removing "
+        "these tests or use an alternative way to generate quantized weights"
+    )
+)
 def test_externalize_rms_norm_with_quantized_linears_ir() -> None:
     """IR check: Quantized (int4) weights retain si4 dtype when externalization re-exports the model."""
 
@@ -3065,7 +3067,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
     )
     quantizer = PostTrainingQuantizer(model, quantization_config)
     model = cast("nn.Module", quantizer.compress())
-    transform_with_custom_compression_ops(model)
+    transform_with_custom_compression_ops(model)  # noqa: F821
 
     sample = (torch.randn(2, 24),)
 
@@ -3109,10 +3111,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 
 @pytest.mark.flaky(reruns=3)
-@pytest.mark.skip(reason=(
-    "transform_with_custom_compression_ops has been deprecated. Consider removing "
-    "these tests or use an alternative way to generate quantized weights"
-))
+@pytest.mark.skip(
+    reason=(
+        "transform_with_custom_compression_ops has been deprecated. Consider removing "
+        "these tests or use an alternative way to generate quantized weights"
+    )
+)
 async def test_externalize_rms_norm_with_quantized_linears() -> None:
     """Quantized (int4) weights retain si4 dtype when externalization re-exports the model.
 
@@ -3153,7 +3157,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
     )
     quantizer = PostTrainingQuantizer(model, quantization_config)
     model = cast("nn.Module", quantizer.compress())
-    transform_with_custom_compression_ops(model)
+    transform_with_custom_compression_ops(model)  # noqa: F821
 
     sample = (torch.randn(2, 24),)
 
@@ -3176,10 +3180,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 @pytest.mark.ir
 @pytest.mark.flaky(reruns=3)
-@pytest.mark.skip(reason=(
-    "transform_with_custom_compression_ops has been deprecated. Consider removing "
-    "these tests or use an alternative way to generate quantized weights"
-))
+@pytest.mark.skip(
+    reason=(
+        "transform_with_custom_compression_ops has been deprecated. Consider removing "
+        "these tests or use an alternative way to generate quantized weights"
+    )
+)
 def test_externalize_gather_mm_with_quantized_rhs_ir() -> None:
     """IR check: Quantized expert weight flows as rhs into an externalized GatherMM composite."""
     num_experts = 4
@@ -3217,7 +3223,7 @@ def forward(self, x: torch.Tensor, indices: torch.Tensor) -> torch.Tensor:
     )
     quantizer = PostTrainingQuantizer(model, quantization_config)
     model = cast("nn.Module", quantizer.compress())
-    transform_with_custom_compression_ops(model)
+    transform_with_custom_compression_ops(model)  # noqa: F821
 
     x = torch.randn(2, 1, 1, in_dim)
     indices = torch.tensor([[0, 2], [1, 3]], dtype=torch.int16)
@@ -3263,10 +3269,12 @@ def forward(self, x: torch.Tensor, indices: torch.Tensor) -> torch.Tensor:
 
 
 @pytest.mark.flaky(reruns=3)
-@pytest.mark.skip(reason=(
-    "transform_with_custom_compression_ops has been deprecated. Consider removing "
-    "these tests or use an alternative way to generate quantized weights"
-))
+@pytest.mark.skip(
+    reason=(
+        "transform_with_custom_compression_ops has been deprecated. Consider removing "
+        "these tests or use an alternative way to generate quantized weights"
+    )
+)
 async def test_externalize_gather_mm_with_quantized_rhs() -> None:
     """Quantized expert weight flows as rhs into an externalized GatherMM composite.
 
@@ -3318,7 +3326,7 @@ def forward(self, x: torch.Tensor, indices: torch.Tensor) -> torch.Tensor:
     )
     quantizer = PostTrainingQuantizer(model, quantization_config)
     model = cast("nn.Module", quantizer.compress())
-    transform_with_custom_compression_ops(model)
+    transform_with_custom_compression_ops(model)  # noqa: F821
 
     x = torch.randn(2, 1, 1, in_dim)
     indices = torch.tensor([[0, 2], [1, 3]], dtype=torch.int16)
@@ -3345,10 +3353,12 @@ def forward(self, x: torch.Tensor, indices: torch.Tensor) -> torch.Tensor:
 
 @pytest.mark.ir
 @pytest.mark.flaky(reruns=3)
-@pytest.mark.skip(reason=(
-    "transform_with_custom_compression_ops has been deprecated. Consider removing "
-    "these tests or use an alternative way to generate quantized weights"
-))
+@pytest.mark.skip(
+    reason=(
+        "transform_with_custom_compression_ops has been deprecated. Consider removing "
+        "these tests or use an alternative way to generate quantized weights"
+    )
+)
 def test_externalize_multiple_composites_with_quantized_weights_ir() -> None:
     """IR check: Multiple composite ops (RMSNorm + SDPA) externalized with quantized linears."""
     head_dim = 16
@@ -3390,7 +3400,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
     )
     quantizer = PostTrainingQuantizer(model, quantization_config)
     model = cast("nn.Module", quantizer.compress())
-    transform_with_custom_compression_ops(model)
+    transform_with_custom_compression_ops(model)  # noqa: F821
 
     sample = (torch.randn(1, 4, embed_dim),)
 
@@ -3433,10 +3443,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 
 @pytest.mark.flaky(reruns=3)
-@pytest.mark.skip(reason=(
-    "transform_with_custom_compression_ops has been deprecated. Consider removing "
-    "these tests or use an alternative way to generate quantized weights"
-))
+@pytest.mark.skip(
+    reason=(
+        "transform_with_custom_compression_ops has been deprecated. Consider removing "
+        "these tests or use an alternative way to generate quantized weights"
+    )
+)
 async def test_externalize_multiple_composites_with_quantized_weights() -> None:
     """Multiple composite ops (RMSNorm + SDPA) externalized with quantized linears.
 
@@ -3484,7 +3496,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
     )
     quantizer = PostTrainingQuantizer(model, quantization_config)
     model = cast("nn.Module", quantizer.compress())
-    transform_with_custom_compression_ops(model)
+    transform_with_custom_compression_ops(model)  # noqa: F821
 
     sample = (torch.randn(1, 4, embed_dim),)
 
@@ -3512,10 +3524,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 
 @pytest.mark.ir
-@pytest.mark.skip(reason=(
-    "transform_with_custom_compression_ops has been deprecated. Consider removing "
-    "these tests or use an alternative way to generate quantized weights"
-))
+@pytest.mark.skip(
+    reason=(
+        "transform_with_custom_compression_ops has been deprecated. Consider removing "
+        "these tests or use an alternative way to generate quantized weights"
+    )
+)
 def test_externalize_gather_mm_combined_with_rms_norm_ir() -> None:
     """IR check: GatherMM + RMSNorm both externalized alongside quantized weights."""
     num_experts = 4
@@ -3553,7 +3567,7 @@ def forward(self, x: torch.Tensor, indices: torch.Tensor) -> torch.Tensor:
     )
     quantizer = PostTrainingQuantizer(model, quantization_config)
     model = cast("nn.Module", quantizer.compress())
-    transform_with_custom_compression_ops(model)
+    transform_with_custom_compression_ops(model)  # noqa: F821
 
     x = torch.randn(2, in_dim)
     indices = torch.tensor([[0, 2], [1, 3]], dtype=torch.int16)
@@ -3590,10 +3604,13 @@ def forward(self, x: torch.Tensor, indices: torch.Tensor) -> torch.Tensor:
     """
     filecheck_pattern(ir, check_file=pattern)
 
-@pytest.mark.skip(reason=(
-    "transform_with_custom_compression_ops has been deprecated. Consider removing "
-    "these tests or use an alternative way to generate quantized weights"
-))
+
+@pytest.mark.skip(
+    reason=(
+        "transform_with_custom_compression_ops has been deprecated. Consider removing "
+        "these tests or use an alternative way to generate quantized weights"
+    )
+)
 async def test_externalize_gather_mm_combined_with_rms_norm() -> None:
     """GatherMM + RMSNorm both externalized alongside quantized weights.
 
@@ -3637,7 +3654,7 @@ def forward(self, x: torch.Tensor, indices: torch.Tensor) -> torch.Tensor:
     )
     quantizer = PostTrainingQuantizer(model, quantization_config)
     model = cast("nn.Module", quantizer.compress())
-    transform_with_custom_compression_ops(model)
+    transform_with_custom_compression_ops(model)  # noqa: F821
 
     x = torch.randn(2, in_dim)
     indices = torch.tensor([[0, 2], [1, 3]], dtype=torch.int16)
@@ -3666,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)