Repository: stanford-crfm/haliax (language: python)

Issue:
named arrays and `eqxi.while_loop(mode=checkpointed)` do not get a long
the shape checks cause problems


Discussion / hints:

Fix patch:
diff --git a/src/haliax/core.py b/src/haliax/core.py
index 6cf5eca..51b1b89 100644
--- a/src/haliax/core.py
+++ b/src/haliax/core.py
@@ -128,7 +128,10 @@ def tree_flatten(self) -> Any:
     @classmethod
     def tree_unflatten(cls, aux, tree: Any) -> Any:
         assert len(tree) == 1
-        return cls(tree[0], axes=aux)
+        # We don't want check shapes b/c there are intermediate states where the shape is wrong
+        # e.g. in eqxi.while_loop
+        with enable_shape_checks(False):
+            return cls(tree[0], axes=aux)
 
     def has_axis(self, axis: AxisSelection) -> bool:
         """Returns true if the given axis is present in this NamedArray."""


Test patch:
diff --git a/tests/core_test.py b/tests/core_test.py
index eac2958..00e4c26 100644
--- a/tests/core_test.py
+++ b/tests/core_test.py
@@ -795,3 +795,28 @@ class TestModule(eqx.Module):
     mod = TestModule(named1)
 
     assert str(mod).startswith("TestModule(named1=Named(int32{'H': 10, 'W': 20}))")
+
+
+def test_named_arrays_work_in_eqxi_while_loop():
+    H = Axis("H", 10)
+    W = Axis("W", 20)
+
+    named1 = hax.random.uniform(PRNGKey(0), (H, W))
+
+    import equinox.internal as eqxi
+
+    def body_fun(t):
+        i, named1 = t
+        return i + 1, named1 + named1
+
+    def cond_fun(t):
+        i, named1 = t
+        return i < 10
+
+    def loss_fun(named1):
+        i, named1 = eqxi.while_loop(cond_fun, body_fun, (0, named1), kind="checkpointed", max_steps=10)
+        return named1.sum().scalar()
+
+    grad_fun = eqx.filter_value_and_grad(loss_fun)
+
+    grad_fun(named1)