Free array iter bugfix (#307)

* free array iter bugfix

* same thing in torch handler

Author: jjcmoon

effectful/handlers/jax/_terms.py

@@ -100,6 +100,9 @@ def shape(self) -> Expr[tuple[int, ...]]:
     def size(self) -> Expr[int]:
         return jnp.size(cast(jax.Array, self))
 
+    def __len__(self):
+        return self.shape[0]
+
     @property
     def ndim(self) -> Expr[int]:
         return jnp.ndim(cast(jax.Array, self))
@@ -217,15 +220,15 @@ def at(self) -> _IndexUpdateHelper:
         """Return an IndexUpdateHelper for array updates."""
         return _IndexUpdateHelper(self)
 
+    def __iter__(self):
+        raise TypeError("A free array is not iterable.")
+
 
 class _EagerArrayTerm(_ArrayTerm):
     def __init__(self, op, tensor, key):
         new_shape, new_key = _desugar_tensor_index(tensor.shape, key)
         super().__init__(op, jnp.reshape(tensor, new_shape), new_key)
 
-    def __len__(self):
-        return self.shape[0]
-
     def __iter__(self):
         for i in range(len(self)):
             yield self[i]

effectful/handlers/torch.py

@@ -489,6 +489,9 @@ def __matmul__(self, other: torch.Tensor) -> torch.Tensor:
     def __rmatmul__(self, other: torch.Tensor) -> torch.Tensor:
         return torch.matmul(other, typing.cast(torch.Tensor, self))
 
+    def __iter__(self):
+        raise TypeError("A free tensor is not iterable.")
+
 
 @Term.register
 class _EagerTensorTerm(torch.Tensor):

tests/test_handlers_jax.py

@@ -1,4 +1,5 @@
 import jax
+import pytest
 
 import effectful.handlers.jax.numpy as jnp
 from effectful.handlers.jax import bind_dims, jax_getitem, jit, sizesof
@@ -292,6 +293,9 @@ def test_jax_at_updates():
 
 def test_jax_len():
     i = defop(jax.Array, name="i")
+    with pytest.raises(TypeError):
+        len(i())
+
     t = jnp.ones((2, 3, 4))
     t_i = jax_getitem(t, [i()])
     assert len(t_i) == 3
@@ -347,3 +351,9 @@ def test_jax_dimension_addition():
     y2 = jax_getitem(y, (i2(), slice(None), None))
     assert y2.shape == (4, 1)
     assert fvsof(y2) >= {i, i2}
+
+
+def test_jax_iter():
+    i = defop(jax.Array, name="i")
+    with pytest.raises(TypeError):
+        tuple(i())

tests/test_handlers_torch.py

@@ -622,3 +622,11 @@ def test_longtensor_index_variables():
 
     z = x[i()] + y[j()]
     assert isinstance(z, torch.Tensor)
+
+
+def test_tensor_iter():
+    i = defop(torch.Tensor)
+    with pytest.raises(TypeError):
+        len(i())
+    with pytest.raises(TypeError):
+        tuple(i())