Include scalar coupling parameters in Hub size indexing

Store the 16 complex scalar coupling values as length-one tensors, recording their names in SafeTensors metadata. Restore the exact original scalar shapes in the loader. This avoids the Hub counter omitting shape=[] parameters. All 296 state tensors again passed byte-for-byte equality and strict native loading on CPU, without CUDA or inference.

Co-Authored-By: Codex <noreply@openai.com>

README.md

@@ -125,6 +125,8 @@ prompt_ids = tokenizer.encode("The nature of consciousness")
 
 Both export files are necessary. The loader restores all learned tensors **and the saved fixed matrices**, strips the `_orig_mod.` prefix when loading the original compiled-training checkpoint, and requires an exact match to the native model's state. It does not regenerate the trained checkpoint's mixer state, retrain the model, alter its tokenizer, or compile it automatically.
 
+For Hub size indexing, the 16 scalar coupling parameters are stored as one-element tensors in SafeTensors. Their names are recorded in the file metadata, and the loader restores their original scalar shapes before loading the model. This changes only the storage representation, not the model or its values.
+
 To use the original checkpoint instead, download `colm_best_Final.pt` too and call `load_colm(".", weights_format="pt")`. See the [GitHub repository](https://github.com/Eden-Eldith/COLM) for full training, generation, and evaluation scripts.
 
 ### Export Verification

colm_loader.py

@@ -7,6 +7,7 @@ from types import SimpleNamespace
 
 import numpy as np
 import torch
+from safetensors import safe_open
 from safetensors.torch import load_file
 
 from model import COLM
@@ -28,6 +29,14 @@ def safetensors_state(directory):
     """Both files are required: learned parameters plus saved fixed mixers."""
     directory = Path(directory)
     state = load_file(str(directory / "model.safetensors"), device="cpu")
+    with safe_open(str(directory / "model.safetensors"), framework="pt", device="cpu") as archive:
+        metadata = archive.metadata() or {}
+    # The Hub counter omits shape=[] tensors. Store these as shape=[1], then
+    # restore the original scalar shape before loading the unchanged model.
+    for name in json.loads(metadata.get("scalar_parameter_names", "[]")):
+        if state[name].shape != (1,):
+            raise ValueError(f"Expected a stored length-one scalar: {name}")
+        state[name] = state[name].reshape(())
     with np.load(directory / "model_buffers.npz", allow_pickle=False) as archive:
         for name in archive.files:
             if name in state:

conversion_report.json

@@ -6,6 +6,24 @@
   "checkpoint_validation_loss": 1.144872358441353,
   "parameter_tensors": 264,
   "parameter_elements": 498214,
+  "scalar_parameters_stored_as_length_one": [
+    "blocks.0.coupling_alpha",
+    "blocks.1.coupling_alpha",
+    "blocks.10.coupling_alpha",
+    "blocks.11.coupling_alpha",
+    "blocks.12.coupling_alpha",
+    "blocks.13.coupling_alpha",
+    "blocks.14.coupling_alpha",
+    "blocks.15.coupling_alpha",
+    "blocks.2.coupling_alpha",
+    "blocks.3.coupling_alpha",
+    "blocks.4.coupling_alpha",
+    "blocks.5.coupling_alpha",
+    "blocks.6.coupling_alpha",
+    "blocks.7.coupling_alpha",
+    "blocks.8.coupling_alpha",
+    "blocks.9.coupling_alpha"
+  ],
   "parameter_elements_by_dtype": {
     "torch.complex64": 52504,
     "torch.float32": 445710
@@ -26,13 +44,13 @@
     "numpy": "2.2.6"
   },
   "sha256": {
-    "model.safetensors": "def65a16eb57a7365ccee0d9e650472d64fe4abad1c2788848f4d79933ea61f3",
+    "model.safetensors": "01473b0a6b1b595315552db82f29987f50f7db362425e39ff06f0f4b51fbab80",
     "model_buffers.npz": "2715e51c179588c3d74f7c03eec36c4a4b5f0ca658575f6c9e16f406dbcdc815",
     "model.py": "22066752391c65e4f360402fe8a0064107fefd7512d1c568004d68c621553f4c",
     "colm_config.json": "fba7681030b2421453f9d7aaa788384e3ffaebaf68aea200bc1c8e1480fa001c",
     "colm_tokenizer.json": "c28592e50ac0fb8c3094c17bb2159bc325a791727af86c8024a2c5b521b1f4c5",
     "tokenizer_train.py": "6d03d16913878e0d6e955bde1cb20a58db3f854e57dfb5b97efb0e7c75c684db",
-    "colm_loader.py": "03da9755ea8d1b98ce2c903c35c2fa658829db5145aab7f4f07b1f94f22a0ff9",
-    "export_safetensors.py": "7d347ee957faf4ff172b2e5baa7f9213de118718abbac2f8e4569bb98c6d58bb"
+    "colm_loader.py": "c6d763efc9bb1aeda99782e966a93900997bee534cc78e43156003ecf8cba264",
+    "export_safetensors.py": "4a22a7b9683a1ae669414b7fa9650da558bcbd5db908e60e9d700084ca35cb54"
   }
 }

export_safetensors.py

@@ -65,12 +65,18 @@ def main():
         raise AssertionError("State does not partition into parameters and buffers")
     parameters = {n: state[n].contiguous() for n in sorted(parameter_names)}
     buffers = {n: state[n].numpy() for n in sorted(buffer_names)}
+    scalar_names = [n for n, value in parameters.items() if value.ndim == 0]
+    serialized_parameters = {
+        n: value.reshape(1) if n in scalar_names else value
+        for n, value in parameters.items()
+    }
 
-    save_file(parameters, str(directory / "model.safetensors"), metadata={
+    save_file(serialized_parameters, str(directory / "model.safetensors"), metadata={
         "format": "pt",
         "source_checkpoint": args.checkpoint.name,
         "source_sha256": original_hash,
         "fixed_buffers": "model_buffers.npz",
+        "scalar_parameter_names": json.dumps(scalar_names),
         "description": "Native F32/C64 trainable parameters; fixed mixers stored separately.",
     })
     np.savez(directory / "model_buffers.npz", **buffers)
@@ -97,6 +103,7 @@ def main():
         "checkpoint_validation_loss": config["checkpoint"]["validation_loss"],
         "parameter_tensors": len(parameters),
         "parameter_elements": sum(counts.values()),
+        "scalar_parameters_stored_as_length_one": scalar_names,
         "parameter_elements_by_dtype": dict(counts),
         "real_scalar_degrees_of_freedom": sum(
             x.numel() * (2 if x.is_complex() else 1) for x in parameters.values()),

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:def65a16eb57a7365ccee0d9e650472d64fe4abad1c2788848f4d79933ea61f3
-size 2226088
+oid sha256:01473b0a6b1b595315552db82f29987f50f7db362425e39ff06f0f4b51fbab80
+size 2226600