haliax 1.4.dev341__tar.gz → 1.4.dev342__tar.gz

{haliax-1.4.dev341 → haliax-1.4.dev342}/PKG-INFO

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

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

{haliax-1.4.dev341 → haliax-1.4.dev342}/src/haliax/_src/state_dict.py

@@ -58,7 +58,7 @@ def flatten_modules_for_export(t: T) -> T:
     def _flatten_module(module):
         if isinstance(module, ModuleWithStateDictSerialization):
             module = module.flatten_for_export()
-            module = jax.tree.map(
+            module = scan_aware_tree_map(
                 _flatten_module,
                 module,
                 is_leaf=lambda x: x is not module and isinstance(x, ModuleWithStateDictSerialization),
@@ -76,7 +76,7 @@ def unflatten_modules_from_export(t: T, template: T) -> T:
     def _unflatten_module(module, template):
         if isinstance(module, ModuleWithStateDictSerialization):
             module = module.unflatten_from_export(template)
-            module = jax.tree.map(
+            module = scan_aware_tree_map(
                 _unflatten_module,
                 module,
                 template,

{haliax-1.4.dev341 → haliax-1.4.dev342}/tests/test_state_dict.py

@@ -7,8 +7,9 @@ import jax.numpy as jnp
 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 _stack_state_dict, _unstack_state_dict
+from haliax.nn.scan import Stacked, _stack_state_dict, _unstack_state_dict
 from haliax.state_dict import from_state_dict, to_state_dict
 
 
@@ -151,3 +152,26 @@ def test_export_layer_norm():
     new_layer_norm = flat_layer_norm.unflatten_from_export(layer_norm2)
 
     assert layer_norm == new_layer_norm
+
+
+def test_stacked_layer_norm():
+    L = hax.Axis("L", 4)
+    D = hax.Axis("D", 10)
+    E = hax.Axis("E", 20)
+
+    norms = Stacked.init(L, hax.nn.LayerNorm)((D, E), eps=1e-5, use_weight=True, use_bias=True)
+
+    norms_flat = flatten_modules_for_export(norms)
+
+    flat_state_dict = to_state_dict(norms_flat)
+
+    assert flat_state_dict["0.weight"].shape == (D.size * E.size,)
+    assert flat_state_dict["0.bias"].shape == (D.size * E.size,)
+    assert flat_state_dict["1.weight"].shape == (D.size * E.size,)
+
+    # now unflatten it
+    norms2 = Stacked.init(L, hax.nn.LayerNorm)((D, E), eps=1e-5, use_weight=True, use_bias=True)
+
+    new_norms = unflatten_modules_from_export(norms_flat, norms2)
+
+    assert norms == new_norms

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

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