Implement plated einsum

effectful/handlers/jax/_handlers.py

@@ -3,13 +3,14 @@
 import typing
 from collections.abc import Callable, Iterable, Mapping, Sequence
 from types import EllipsisType
-from typing import Annotated
+from typing import Annotated, Any
 
 from opt_einsum import get_symbol
 from opt_einsum.parser import parse_einsum_input
 
 try:
     import jax
+    import jax.core
     import jax.numpy as jnp
 except ImportError:
     raise ImportError("JAX is required to use effectful.handlers.jax")
@@ -31,6 +32,9 @@
 IndexElement = None | int | slice | Sequence[int] | EllipsisType | jax.Array
 
 
+class JaxOperation[**P, T](Operation[P, T]): ...
+
+
 def is_eager_array(x):
     return isinstance(x, jax.Array) or (
         isinstance(x, Term)
@@ -141,7 +145,7 @@ def _register_jax_op[**P, T](jax_fn: Callable[P, T]):
     if getattr(jax_fn, "__name__", None) == "__getitem__":
         return jax_getitem
 
-    @defop
+    @JaxOperation.define
     def _jax_op(*args, **kwargs) -> jax.Array:
         tm = defdata(_jax_op, *args, **kwargs)
         sized_fvs = sizesof(tm)
@@ -158,6 +162,8 @@ def _jax_op(*args, **kwargs) -> jax.Array:
             )
         ):
             raise NotHandled
+        elif _jax_op is jax_getitem and not args[1]:
+            return args[0]
         elif sized_fvs and set(sized_fvs.keys()) == fvsof(tm) - {jax_getitem, _jax_op}:
             # note: this cast is a lie. partial_eval can return non-arrays, as
             # can jax_fn. for example, some jax functions return tuples,
@@ -296,7 +302,7 @@ def jax_getitem(x: jax.Array, key: tuple[IndexElement, ...]) -> jax.Array:
 def bind_dims[T, A, B](
     __dispatch: Callable[[type], Callable[..., T]],
     value: Annotated[T, Scoped[A | B]],
-    *names: Annotated[Operation[[], jax.Array], Scoped[B]],
+    *names: Annotated[Operation[[], Any], Scoped[B]],
 ) -> Annotated[T, Scoped[A]]:
     """Convert named dimensions to positional dimensions.
 
@@ -315,26 +321,40 @@ def bind_dims[T, A, B](
     >>> bind_dims(t, b, a).shape
     (3, 2)
     """
+    if not names:
+        return value
     if isinstance(value, Term) and value.op == bind_dims:
         return bind_dims(value.args[0], *(names + tuple(value.args[1:])))
     if jax.tree_util.treedef_is_leaf(jax.tree.structure(value)):
         return __dispatch(typeof(value))(value, *names)
     return jax.tree.map(lambda v: bind_dims(v, *names), value)
 
 
+@bind_dims.register(object)
+def _(*args, **kwargs):
+    raise NotHandled
+
+
 @defop
 @_CustomSingleDispatchCallable
 def unbind_dims[T, A, B](
     __dispatch: Callable[[type], Callable[..., T]],
     value: Annotated[T, Scoped[A | B]],
-    *names: Annotated[Operation[[], jax.Array], Scoped[B]],
+    *names: Annotated[Operation[[], Any], Scoped[B]],
 ) -> Annotated[T, Scoped[A | B]]:
     """Convert positional dimensions to named dimensions."""
+    if not names:
+        return value
     if jax.tree_util.treedef_is_leaf(jax.tree.structure(value)):
         return __dispatch(typeof(value))(value, *names)
     return jax.tree.map(lambda v: unbind_dims(v, *names), value)
 
 
+@unbind_dims.register(object)
+def _(*args, **kwargs):
+    raise NotHandled
+
+
 def jit(f, *args, **kwargs):
     f_noindex, f_reindex = _indexed_func_wrapper(f, jax_getitem, sizesof)
     f_noindex_jitted = jax.jit(f_noindex, *args, **kwargs)
@@ -388,7 +408,18 @@ def _(x: jax.Array, other) -> bool:
     )
 
 
-@syntactic_hash.register(jax.Array)
+@syntactic_eq.register
+def _(x: jax.core.Tracer, other) -> bool:
+    return isinstance(other, jax.core.Tracer) and id(x) == id(other)
+
+
+@syntactic_hash.register
 def _(x: jax.Array) -> int:
     # Concrete arrays aren't hashable; hash by shape, dtype, and bytes.
     return hash(("jax.Array", x.shape, str(x.dtype), bytes(jax.numpy.asarray(x))))
+
+
+@syntactic_hash.register
+def _(x: jax.core.Tracer) -> int:
+    # Concrete arrays aren't hashable; hash by shape, dtype, and bytes.
+    return hash(("jax.core.Tracer", id(x)))

effectful/handlers/jax/_terms.py

@@ -10,6 +10,7 @@
     IndexElement,
     _register_jax_op,
     bind_dims,
+    is_eager_array,
     jax_getitem,
     unbind_dims,
 )
@@ -461,7 +462,7 @@ def _bind_dims_array(t: jax.Array, *args: Operation[[], jax.Array]) -> jax.Array
         return t
 
     # ensure that the result is a jax_getitem with an array as the first argument
-    if not (t.op is jax_getitem and isinstance(t.args[0], jax.Array)):
+    if not (t.op is jax_getitem and is_eager_array(t)):
         raise NotHandled
 
     array = t.args[0]