haliax 1.4.dev348__tar.gz → 1.4.dev351__tar.gz

{haliax-1.4.dev348 → haliax-1.4.dev351}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev348
+Version: 1.4.dev351
 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.dev351/src/haliax/__about__.py

@@ -0,0 +1 @@
+__version__ = "1.4.dev351"

{haliax-1.4.dev348 → haliax-1.4.dev351}/src/haliax/_src/state_dict.py

@@ -215,6 +215,8 @@ def from_state_dict(tree: T, state_dict: StateDict, prefix: Optional[str] = None
             raise ValueError("Cannot extract a leaf value from a state dict without a prefix")
         # TODO: add "strict" flag so we can return None in cases where it's just missing
         return jnp.array(state_dict[prefix])
+    elif tree is None:
+        return None
     else:
         if prefix is None:
             return tree

{haliax-1.4.dev348 → haliax-1.4.dev351}/tests/test_state_dict.py

@@ -175,3 +175,27 @@ def test_stacked_layer_norm():
     new_norms = unflatten_modules_from_export(norms_flat, norms2)
 
     assert norms == new_norms
+
+
+def test_linear_doesnt_read_bias_if_it_didnt_have_bias():
+    H = hax.Axis("H", 10)
+    W = hax.Axis("W", 20)
+    D = hax.Axis("D", 30)
+    B = hax.Axis("B", 40)
+
+    linear = hax.nn.Linear.init((H, W), (D, B), key=jax.random.PRNGKey(0), use_bias=False, out_first=True)
+
+    flat_linear = linear.flatten_for_export()
+
+    flat_state_dict = to_state_dict(flat_linear)
+
+    assert "bias" not in flat_state_dict
+    flat_state_dict["bias"] = jnp.zeros((D.size * B.size,))  # add a dummy bias
+
+    # now unflatten it
+    linear2 = Linear.init((H, W), (D, B), key=jax.random.PRNGKey(1), use_bias=False, out_first=True)
+    flinear2 = linear2.flatten_for_export()
+    flinear2 = from_state_dict(flinear2, flat_state_dict)
+    new_linear = flinear2.unflatten_from_export(linear2)
+
+    assert linear == new_linear

haliax-1.4.dev348/src/haliax/__about__.py

@@ -1 +0,0 @@
-__version__ = "1.4.dev348"