haliax 1.4.dev447__tar.gz → 1.4.dev452__tar.gz

{haliax-1.4.dev447 → haliax-1.4.dev452}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev447
+Version: 1.4.dev452
 Summary: Named Tensors for Legible Deep Learning in JAX
 Project-URL: Homepage, https://github.com/stanford-crfm/haliax
 Project-URL: Bug Tracker, https://github.com/stanford-crfm/haliax/issues/

{haliax-1.4.dev447 → haliax-1.4.dev452}/pyproject.toml

@@ -45,6 +45,7 @@ dev = [
     "pymdown-extensions",
     "pygments",
     "chex>=0.1.86",
+    "safetensors>=0.4.3",
     "pre-commit",
 ]
 

{haliax-1.4.dev447 → haliax-1.4.dev452}/src/haliax/__about__.py

@@ -3,4 +3,4 @@
 # SPDX-License-Identifier: Apache-2.0
 
 
-__version__ = "1.4.dev447"
+__version__ = "1.4.dev452"

{haliax-1.4.dev447 → haliax-1.4.dev452}/src/haliax/_src/state_dict.py

@@ -26,8 +26,10 @@ from haliax.tree_util import scan_aware_tree_map
 
 try:
     import safetensors
+    import safetensors.numpy as safetensors_numpy
 except ImportError:
     safetensors = None
+    safetensors_numpy = None
 
 
 StateDict = dict[str, Any]
@@ -436,7 +438,9 @@ def save_state_dict(state_dict: StateDict, path):
     state_dict = {k: v for k, v in state_dict.items() if v is not None}
     if jax.process_index() == 0:
         # the "pt" is a lie but it doesn't seem to actually matter and HF demands it
-        safetensors.numpy.save_file(state_dict, path, metadata={"format": "pt"})
+        if safetensors_numpy is None:
+            raise ImportError("safetensors_numpy is not installed")
+        safetensors_numpy.save_file(state_dict, path, metadata={"format": "pt"})
     global _GLOBAL_SAVE_COUNT
     sync_global_devices(f"save_state_dict {_GLOBAL_SAVE_COUNT}")
     _GLOBAL_SAVE_COUNT += 1
@@ -447,5 +451,7 @@ def load_state_dict(path):
     Load a model's state dict from a file, bringing all tensors to the CPU first and then converting to numpy.
     This will load using safetensors format
     """
-    state_dict = safetensors.numpy.load_file(path)
+    if safetensors_numpy is None:
+        raise ImportError("safetensors_numpy is not installed")
+    state_dict = safetensors_numpy.load_file(path)
     return state_dict

{haliax-1.4.dev447 → haliax-1.4.dev452}/src/haliax/nn/linear.py

@@ -158,12 +158,17 @@ class Linear(ModuleWithStateDictSerialization, ReparamEnabled):
             return self.weight.axes[-len(self.Out) :] != self.Out
 
     def to_state_dict(self, prefix: Optional[str] = None) -> StateDict:
-        scaled = dataclasses.replace(self, weight=self.weight * self.reparam.active_scale)
+        # weight can be None for certain filtering things like LoRA
+        scaled = dataclasses.replace(
+            self, weight=self.weight * self.reparam.active_scale if self.weight is not None else None
+        )
         return default_eqx_module_to_state_dict(scaled, prefix)
 
     def from_state_dict(self: Mod, state_dict: StateDict, prefix: Optional[str] = None) -> Mod:
         unscaled = default_eqx_module_from_state_dict(self, state_dict, prefix)
-        return dataclasses.replace(unscaled, weight=unscaled.weight / self.reparam.active_scale)
+        if unscaled.weight is not None:
+            unscaled = dataclasses.replace(unscaled, weight=unscaled.weight / self.reparam.active_scale)
+        return unscaled
 
     @staticmethod
     def input_reparam(use_mup: bool = True) -> type[AbstractLinearReparam]:

{haliax-1.4.dev447 → haliax-1.4.dev452}/tests/test_state_dict.py

@@ -9,13 +9,14 @@ from typing import Any
 import equinox as eqx
 import jax
 import jax.numpy as jnp
+import numpy as np
 import pytest
 
 import haliax as hax
 from haliax._src.state_dict import flatten_modules_for_export, unflatten_modules_from_export
 from haliax.nn import Linear
 from haliax.nn.scan import Stacked, _stack_state_dict, _unstack_state_dict
-from haliax.state_dict import from_state_dict, to_state_dict
+from haliax.state_dict import from_state_dict, load_state_dict, save_state_dict, to_state_dict
 
 
 @pytest.mark.parametrize("out_dims_first", [True, False])
@@ -159,6 +160,44 @@ def test_export_layer_norm():
     assert layer_norm == new_layer_norm
 
 
+def test_save_state_dict_roundtrip(tmp_path):
+    pytest.importorskip("safetensors.numpy")
+    state_dict = {
+        "weight": np.arange(6, dtype=np.float32).reshape(2, 3),
+        "bias": np.linspace(0.0, 1.0, 3, dtype=np.float32),
+    }
+    path = tmp_path / "state.safetensors"
+
+    save_state_dict(state_dict, path)
+
+    assert path.exists()
+    loaded = load_state_dict(path)
+
+    assert set(loaded.keys()) == {"weight", "bias"}
+    np.testing.assert_array_equal(loaded["weight"], state_dict["weight"])
+    np.testing.assert_array_equal(loaded["bias"], state_dict["bias"])
+
+
+def test_save_state_dict_filters_none_and_sets_metadata(tmp_path):
+    safetensors = pytest.importorskip("safetensors")
+
+    state_dict = {
+        "weight": np.ones((4,), dtype=np.float32),
+        "optional": None,
+    }
+    path = tmp_path / "state.safetensors"
+
+    save_state_dict(state_dict, path)
+
+    with safetensors.safe_open(str(path), framework="numpy") as f:
+        keys = list(f.keys())
+
+        assert "weight" in keys
+        assert "optional" not in keys
+        np.testing.assert_array_equal(f.get_tensor("weight"), state_dict["weight"])
+        assert f.metadata().get("format") == "pt"
+
+
 def test_stacked_layer_norm():
     L = hax.Axis("L", 4)
     D = hax.Axis("D", 10)

{haliax-1.4.dev447 → haliax-1.4.dev452}/uv.lock

@@ -315,6 +315,7 @@ dev = [
     { name = "pygments" },
     { name = "pymdown-extensions" },
     { name = "pytest" },
+    { name = "safetensors" },
 ]
 
 [package.metadata]
@@ -342,6 +343,7 @@ dev = [
     { name = "pygments" },
     { name = "pymdown-extensions" },
     { name = "pytest", specifier = ">=7.4.0" },
+    { name = "safetensors", specifier = ">=0.4.3" },
 ]
 
 [[package]]