Fix loading of palettizer state dict with no indices into prepared model

src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py

@@ -257,15 +257,22 @@ def _save_to_state_dict(self, destination, prefix, keep_vars):
         if self._indices_stale:
             destination.pop(prefix + "indices", None)
 
-    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
+    def _load_from_state_dict(
+        self, state_dict, prefix, local_metadata, strict, missing_keys, *args, **kwargs
+    ):
         """Load centroids from a checkpoint, reconstructing them from a legacy
-        ``lut`` buffer when present.
+        ``lut`` buffer when present. Absent ``indices`` are not reported as missing;
+        they are recomputed from centroids on next use.
         """
         lut_key, centroids_key, index_key = prefix + "lut", prefix + "centroids", prefix + "indices"
         if centroids_key not in state_dict and lut_key in state_dict:
             state_dict[centroids_key] = self._centroids_from_lut(state_dict[lut_key])
         had_indices = index_key in state_dict
-        super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
+        super()._load_from_state_dict(
+            state_dict, prefix, local_metadata, strict, missing_keys, *args, **kwargs
+        )
+        if not had_indices and index_key in missing_keys:
+            missing_keys.remove(index_key)
         if self.centroids is not None:
             self._centroids_initialized = True
             scale_ok = not self.enable_per_channel_scale or self.per_channel_scale is not None

tests/palettization/test_kmeans_fake_palettize.py

@@ -1884,6 +1884,23 @@ def test_load_from_state_dict_with_and_without_indices(self):
         assert target._centroids_initialized is True
         assert target._indices_stale is True
 
+    def test_strict_load_without_indices_into_initialized_module(self):
+        """A prepared target loads an indices-less state dict under strict and recomputes
+        indices.
+        """
+        spec = PalettizationSpec(n_bits=2, granularity=PerTensorGranularity())
+        weight = torch.randn(8, 8)
+        source = _KMeansFakePalettize(**spec.__dict__)
+        source._initialize(weight)
+        state_dict = source.state_dict()
+        del state_dict["indices"]
+
+        target = _KMeansFakePalettize(**spec.__dict__)
+        target._initialize(torch.randn(8, 8))
+        target.load_state_dict(state_dict)  # strict=True: must not raise
+        assert target._indices_stale is True
+        assert torch.equal(target.hard_assign(weight), source.hard_assign(weight))
+
     def test_load_from_state_dict_missing_scale_is_stale(self):
         """With per-channel scale enabled, loading indices but no scale marks stale
         so both are recomputed together on next use.