From e5bc5bda622baed5e22df71ed91d44ce080758f0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ey=C3=BCp=20Can=20Akman?= <eyupcanakman@gmail.com>
Date: Fri, 26 Jun 2026 13:07:24 +0300
Subject: [PATCH] fix(quantization): export root-module weights in graph CoreML

A weight on the root module (a bare nn.Linear) exports with a dot-less
get_attr target, which _get_weight_input_names rejected with "Invalid
weight target path". Treat a dot-less target as a root-module parameter
(module name ""), which _get_weight_module already resolves.
---
 .../_graph/_prepare_for_export.py             | 10 +---
 .../test_pt2e_weight_only_mil_export.py       | 60 +++++++++++++++++++
 2 files changed, 63 insertions(+), 7 deletions(-)

diff --git a/src/coreai_opt/quantization/_graph/_prepare_for_export.py b/src/coreai_opt/quantization/_graph/_prepare_for_export.py
index 36278d6..e2c37d6 100644
--- a/src/coreai_opt/quantization/_graph/_prepare_for_export.py
+++ b/src/coreai_opt/quantization/_graph/_prepare_for_export.py
@@ -142,12 +142,11 @@ def _get_weight_input_names(
 
     Returns:
         Tuple of (module_name, param_name)
-        - module_name: e.g., "conv", "layer1.0"
+        - module_name: e.g., "conv", "layer1.0", or "" for a root-module parameter
         - param_name: e.g., "weight", "bias"
 
     Raises:
         ValueError: If node is not a weight quantization node
-        ValueError: If weight target path is invalid
 
     """
     if not _is_weight_fake_quant(fake_quant_node, module):
@@ -159,13 +158,10 @@ def _get_weight_input_names(
     # Extract module and parameter name from target path
     # e.g., "conv.weight" -> ("conv", "weight")
     # e.g., "layer1.0.weight" -> ("layer1.0", "weight")
+    # e.g., "weight" -> ("", "weight") for a parameter on the root module
     target_path = str(input_node.target)
     last_dot_idx = target_path.rfind(".")
-    if last_dot_idx == -1:
-        msg = f"Invalid weight target path: {target_path}"
-        raise ValueError(msg)
-
-    module_name = target_path[:last_dot_idx]
+    module_name = target_path[:last_dot_idx] if last_dot_idx != -1 else ""
     param_name = target_path[last_dot_idx + 1 :]
 
     return module_name, param_name
diff --git a/tests/export/test_pt2e_weight_only_mil_export.py b/tests/export/test_pt2e_weight_only_mil_export.py
index a083d09..d1af8d2 100644
--- a/tests/export/test_pt2e_weight_only_mil_export.py
+++ b/tests/export/test_pt2e_weight_only_mil_export.py
@@ -83,3 +83,63 @@ def test_simple_model_weight_only_mil_export(
         export_backend=backend,
         prepared_model_output=prepared_model_output,
     )
+
+
+@pytest.mark.parametrize("dtype", ["int8", "uint8"])
+@pytest.mark.parametrize(
+    "granularity",
+    [
+        PerTensorGranularity(),
+        PerChannelGranularity(axis=0),
+        PerBlockGranularity(axis=0, block_size=2),
+    ],
+)
+def test_root_module_weight_only_mil_export(
+    dtype: str,
+    granularity: PerTensorGranularity | PerChannelGranularity | PerBlockGranularity,
+) -> None:
+    """Test weight-only quantization export to CoreML for a root-module weight.
+
+    A bare ``nn.Linear`` owns its weight directly, so the fake-quant input target is the
+    dot-less ``"weight"`` rather than ``"<submodule>.weight"``. The graph CoreML export
+    path rejected that with ``Invalid weight target path: weight`` instead of resolving it
+    to the root module.
+    """
+    model = torch.nn.Linear(8, 4)
+    model.eval()
+    model_input = torch.randn(2, 8)
+
+    config = QuantizerConfig(
+        global_config=ModuleQuantizerConfig(
+            op_state_spec={
+                "weight": QuantizationSpec(
+                    dtype=dtype,
+                    qscheme=QuantizationScheme.SYMMETRIC,
+                    granularity=granularity,
+                    fake_quantize_cls="default",
+                    qparam_calculator_cls="default",
+                    range_calculator_cls="minmax",
+                ),
+            },
+            op_input_spec=None,
+            op_output_spec=None,
+        ),
+    )
+
+    quantizer = Quantizer(model, config)
+    prepared_model = quantizer.prepare((model_input,))
+
+    with torch.no_grad():
+        prepared_model_output = prepared_model(model_input)
+
+    finalized_model = quantizer.finalize(backend=ExportBackend.CoreML)
+
+    export_utils.convert_and_verify(
+        finalized_model=finalized_model,
+        input_data=model_input,
+        expected_ops={
+            "constexpr_blockwise_shift_scale": 1,
+        },
+        export_backend=ExportBackend.CoreML,
+        prepared_model_output=prepared_model_output,
+    )