Update indexing in Skala 1.1 model for faster backward

src/skala/functional/_hashes.py

@@ -8,19 +8,23 @@
 security-sensitive.
 """
 
-# Maps (repo_id, filename) -> expected SHA-256 hex digest.
-KNOWN_HASHES: dict[tuple[str, str], str] = {
+# Maps (repo_id, filename) -> accepted SHA-256 hex digests.
+KNOWN_HASHES: dict[tuple[str, str], tuple[str, ...]] = {
     ("microsoft/skala-1.0", "skala-1.0.fun"): (
-        "08d94436995937eb57c451af7c92e2c7f9e1bff6b7da029a3887e9f9dd4581c0"
+        "08d94436995937eb57c451af7c92e2c7f9e1bff6b7da029a3887e9f9dd4581c0",
     ),
     ("microsoft/skala-1.0", "skala-1.0-cuda.fun"): (
-        "0b38e13237cec771fed331664aace42f8c0db8f15caca6a5c563085e61e2b1fd"
+        "0b38e13237cec771fed331664aace42f8c0db8f15caca6a5c563085e61e2b1fd",
     ),
     ("microsoft/skala-1.1", "skala-1.1.fun"): (
-        "0c8432ac3f03c8f1276372df9aca5b7ee7f8939d47a8789eb158976e89aa0606"
+        "0c8432ac3f03c8f1276372df9aca5b7ee7f8939d47a8789eb158976e89aa0606",
+        "7f3e8622e1eb520ccd88a55464c3e359ac4d7e5ccbd1fb77a26afa1e1c20a5cd",
     ),
     (
         "microsoft/skala-1.1",
         "skala-1.1-cuda.fun",
-    ): "f77be6002d873c0a2384b6df7850d32bbec519036344ff5fdde9730c6f9a4326",
+    ): (
+        "f77be6002d873c0a2384b6df7850d32bbec519036344ff5fdde9730c6f9a4326",
+        "f848eae769dca91741a518ae7275d10caac398ab21db649f91bc1f136872f223",
+    ),
 }

src/skala/functional/load.py

@@ -8,7 +8,7 @@
 import io
 import json
 import os
-from collections.abc import Iterable, Mapping
+from collections.abc import Collection, Iterable, Mapping
 from typing import IO, Any, cast
 
 import torch
@@ -80,15 +80,16 @@ def load(
         fp: str | os.PathLike[str] | IO[bytes],
         device: torch.device | None = None,
         *,
-        expected_hash: str | None = None,
+        expected_hash: str | Collection[str] | None = None,
     ) -> "TracedFunctional":
         """Load a traced functional from a file.
 
         Args:
             fp: File path or readable binary stream.
             device: Target device for the loaded model.
-            expected_hash: If provided, the SHA-256 hex digest that the file
-                content must match. A ``ValueError`` is raised on mismatch.
+            expected_hash: If provided, a SHA-256 hex digest or collection of
+                accepted digests that the file content must match. A ``ValueError``
+                is raised on mismatch.
         """
         extra_files = {
             "metadata": b"",
@@ -101,6 +102,11 @@ def load(
             device = torch.get_default_device()
 
         if expected_hash is not None:
+            expected_hashes = (
+                (expected_hash,)
+                if isinstance(expected_hash, str)
+                else tuple(expected_hash)
+            )
             # Read the whole file into memory so we can hash it before
             # passing it to the unsafe torch.jit.load deserializer.
             if isinstance(fp, (str, os.PathLike)):
@@ -110,10 +116,15 @@ def load(
                 # Assume a file-like object
                 data = fp.read()
             actual_hash = hashlib.sha256(data).hexdigest()
-            if actual_hash != expected_hash:
+            if actual_hash not in expected_hashes:
+                expected = (
+                    expected_hashes[0]
+                    if len(expected_hashes) == 1
+                    else f"one of {', '.join(expected_hashes)}"
+                )
                 raise ValueError(
                     f"Hash mismatch for functional file: "
-                    f"expected {expected_hash}, got {actual_hash}. "
+                    f"expected {expected}, got {actual_hash}. "
                     f"The file may have been tampered with."
                 )
             fp = io.BytesIO(data)

src/skala/functional/model.py

@@ -736,11 +736,16 @@ def _forward_batched_tp_down(
             xw = xw * distance_weights  # broadcasts over B
         # Flatten instruction & channel dims: (B, s, r, w) -> (s, r, B*w)
         xw_flat = xw.permute(1, 2, 0, 3).reshape(xw.shape[1], xw.shape[2], -1)
-        return (
-            xw_flat[:, :, self._tp_down_xw_idx]
-            * x2[:, :, self._tp_down_sph_idx]
-            * self._tp_down_norm.to(x2.dtype)
+        # Use explicit gather to avoid Inductor miscompile observed with advanced indexing.
+        idx_xw = self._tp_down_xw_idx.view(1, 1, -1).expand(
+            xw_flat.shape[0], xw_flat.shape[1], -1
+        )
+        idx_x2 = self._tp_down_sph_idx.view(1, 1, -1).expand(
+            x2.shape[0], x2.shape[1], -1
         )
+        xw_sel = torch.gather(xw_flat, dim=2, index=idx_xw)
+        x2_sel = torch.gather(x2, dim=2, index=idx_x2)
+        return xw_sel * x2_sel * self._tp_down_norm.to(x2.dtype)
 
     def _forward_batched_tp_up(
         self,

tests/test_hash_pinning.py

@@ -53,6 +53,16 @@ def test_load_with_correct_hash(dummy_fun_bytes: bytes) -> None:
     assert isinstance(func, TracedFunctional)
 
 
+def test_load_with_one_of_multiple_hashes(dummy_fun_bytes: bytes) -> None:
+    """Loading succeeds when any accepted hash matches."""
+    correct_hash = hashlib.sha256(dummy_fun_bytes).hexdigest()
+    func = TracedFunctional.load(
+        io.BytesIO(dummy_fun_bytes),
+        expected_hash=("0" * 64, correct_hash),
+    )
+    assert isinstance(func, TracedFunctional)
+
+
 def test_load_with_wrong_hash(dummy_fun_bytes: bytes) -> None:
     """Loading raises ValueError when the hash does not match."""
     wrong_hash = "0" * 64