Remove buggy check in jax tensor indexing (#302)

* fix advanced indexing bug

* format

Author: jfeser

effectful/handlers/jax/_handlers.py

@@ -77,10 +77,6 @@ def update_sizes(sizes, op, size):
 
     def _getitem_sizeof(x: jax.Array, key: tuple[Expr[IndexElement], ...]):
         if is_eager_array(x):
-            if len(key) > x.ndim:
-                raise IndexError(
-                    f"Indexing with too many dimensions: expected {x.ndim} got {len(key)}"
-                )
             for i, k in enumerate(key):
                 if isinstance(k, Term) and len(k.args) == 0 and len(k.kwargs) == 0:
                     update_sizes(sizes, k.op, x.shape[i])

tests/test_handlers_jax.py

@@ -305,3 +305,52 @@ def test_jax_len():
 
     for row in t_i:
         assert len(row) == 4
+
+
+def test_jax_dimension_addition():
+    """Test jax_getitem with dimension addition via None indexing."""
+    i = defop(jax.Array, name="i")
+    i2 = defop(i)
+
+    # Basic case: indexing with slice and defop
+    x = jax_getitem(jnp.eye(3), (i(), slice(None)))
+    assert x.shape == (3,)
+    assert fvsof(x) >= {i}
+
+    # Multiple defops with None - this should work fine
+    x2 = jax_getitem(jnp.eye(3), (i(), i2(), None))
+    assert x2.shape == (1,)
+    assert fvsof(x2) >= {i, i2}
+
+    # The problematic case: indexing a Term with defop and None
+    # This should work but may currently fail
+    x3 = jax_getitem(x, (i2(), None))
+    assert x3.shape == (1,)
+    assert fvsof(x3) >= {i, i2}
+
+    # Additional test cases for dimension addition
+    # Test with multiple None dimensions
+    x4 = jax_getitem(x, (i2(), None, None))
+    assert x4.shape == (1, 1)
+    assert fvsof(x4) >= {i, i2}
+
+    # Test with slice and None
+    x5 = jax_getitem(x, (slice(None), None))
+    assert x5.shape == (3, 1)
+    assert fvsof(x5) >= {i}
+
+    # Test with Ellipsis and None
+    x6 = jax_getitem(x, (..., None))
+    assert x6.shape == (3, 1)
+    assert fvsof(x6) >= {i}
+
+    # Test nested indexing with dimension addition
+    base = jnp.ones((2, 3, 4))
+    y = jax_getitem(base, (i(), slice(None), slice(None)))
+    assert y.shape == (3, 4)
+    assert fvsof(y) >= {i}
+
+    # Index the result with another defop and add dimension
+    y2 = jax_getitem(y, (i2(), slice(None), None))
+    assert y2.shape == (4, 1)
+    assert fvsof(y2) >= {i, i2}