BasisResearch
diff --git a/‎effectful/handlers/jax/_handlers.py‎
Lines changed: 108 additions & 6 deletions b/‎effectful/handlers/jax/_handlers.py‎
Lines changed: 108 additions & 6 deletions
diff --git a/‎effectful/handlers/jax/_terms.py‎
Lines changed: 53 additions & 18 deletions b/‎effectful/handlers/jax/_terms.py‎
Lines changed: 53 additions & 18 deletions
@@ -1,17 +1,21 @@
 import functools
+import itertools
 import typing
-from collections.abc import Callable, Mapping, Sequence
+from collections.abc import Callable, Iterable, Mapping, Sequence
 from types import EllipsisType
 from typing import Annotated
 
+from opt_einsum import get_symbol
+from opt_einsum.parser import parse_einsum_input
+
 try:
     import jax
     import jax.numpy as jnp
 except ImportError:
     raise ImportError("JAX is required to use effectful.handlers.jax")
 
 from effectful.internals.runtime import interpreter
-from effectful.ops.semantics import apply, evaluate, fvsof, typeof
+from effectful.ops.semantics import apply, evaluate, fvsof, fwd, typeof
 from effectful.ops.syntax import (
     Scoped,
     _CustomSingleDispatchCallable,
@@ -66,9 +70,13 @@ def update_sizes(sizes, op, size):
 
     def _getitem_sizeof(x: jax.Array, key: tuple[Expr[IndexElement], ...]):
         if is_eager_array(x):
-            for i, k in enumerate(key):
+            i = 0
+            for k in key:
                 if isinstance(k, Term) and len(k.args) == 0 and len(k.kwargs) == 0:
                     update_sizes(sizes, k.op, x.shape[i])
+                if k is not None:
+                    i += 1
+
         return defdata(jax_getitem, x, key)
 
     def _apply(op, *args, **kwargs):
@@ -89,8 +97,9 @@ def _partial_eval(t: Expr[jax.Array]) -> Expr[jax.Array]:
 
     # if any dimension is zero sized, the result is empty
     if any(size == 0 for size in sized_fvs.values()):
-        key = tuple(sized_fvs.keys())
-        shape = tuple(sized_fvs[k] for k in key)
+        ops = tuple(sized_fvs.keys())
+        key = tuple(k() for k in ops)
+        shape = tuple(sized_fvs[k] for k in ops)
         return jax_getitem(jnp.empty(shape), key)
 
     def _is_eager(t):
@@ -142,7 +151,11 @@ def _jax_op(*args, **kwargs) -> jax.Array:
             and not isinstance(args[0], Term)
             and sized_fvs
             and args[1]
-            and all(isinstance(k, Term) and k.op in sized_fvs for k in args[1])
+            and all(
+                (isinstance(k, Term) and k.op in sized_fvs)
+                or (isinstance(k, slice) and k == slice(None))
+                for k in args[1]
+            )
         ):
             raise NotHandled
         elif sized_fvs and set(sized_fvs.keys()) == fvsof(tm) - {jax_getitem, _jax_op}:
@@ -182,6 +195,93 @@ def _jax_op(*args, **kwargs) -> jax.Array:
     return _jax_op
 
 
+def _named_dims(term: Expr[jax.Array]) -> tuple[Operation, ...]:
+    if not (isinstance(term, Term) and term.op == jax_getitem):
+        return ()
+    index = term.args[1]
+    assert isinstance(index, Iterable)
+    return tuple(i.op for i in index if isinstance(i, Term) and not i.args)
+
+
+def _reduce_named(array, axis=None, **kwargs) -> jax.Array:
+    if axis is None:
+        return fwd()
+
+    named_dims = _named_dims(array)
+    if not named_dims:
+        return fwd()
+
+    bound_arr = bind_dims(array, *named_dims)
+
+    if isinstance(axis, int):
+        axis = (axis,)
+    shifted_axis = tuple(a + len(named_dims) if a >= 0 else a for a in axis)
+
+    reduced = fwd(bound_arr, axis=shifted_axis, **kwargs)
+    return unbind_dims(reduced, *named_dims)
+
+
+def _einsum_named(subscripts, *operands, **kwargs) -> jax.Array:
+    # only the string-subscripts form is handled; forward the interleaved form
+    if not isinstance(subscripts, str):
+        if any(isinstance(x, Term) for x in (subscripts, *operands)):
+            raise ValueError("Interleaved einsum is not implemented with named tensors")
+        return jax.numpy.einsum(subscripts, *operands, **kwargs)
+
+    # forward if any operand has a symbolic (Term) shape
+    if any(isinstance(arr.shape, Term) for arr in operands):
+        raise NotHandled
+
+    named = [_named_dims(op) for op in operands]
+
+    # normalize: expand ellipses and make the output explicit, using the
+    # positional shapes (shapes=True avoids materializing the operands)
+    shapes = [op.shape for op in operands]
+    in_part, out_part, _ = parse_einsum_input([subscripts, *shapes], shapes=True)
+    in_specs = in_part.split(",")
+    assert len(in_specs) == len(operands)
+
+    # fresh symbols for named dims, avoiding every symbol already in use;
+    # get_symbol gives an effectively unlimited supply (spills into unicode)
+    used = {c for c in (in_part + out_part) if c not in ",->"}
+    counter = itertools.count()
+
+    def next_symbol():
+        while True:
+            s = get_symbol(next(counter))
+            if s not in used:
+                used.add(s)
+                return s
+
+    # assign a letter per unique named dim; shared names reuse the same letter
+    # so einsum aligns them as batch dims rather than contracting
+    letter_of, order = {}, []
+    for dims in named:
+        for d in dims:
+            if d not in letter_of:
+                letter_of[d] = next_symbol()
+                order.append(d)
+
+    # bind named dims to leading positional axes and prepend their letters
+    bound, new_in_specs = [], []
+    for op, dims, spec in zip(operands, named, in_specs):
+        bound.append(bind_dims(op, *dims) if dims else op)
+        new_in_specs.append("".join(letter_of[d] for d in dims) + spec)
+
+    # add every named dim to the front of the output as passthrough
+    out_prefix = "".join(letter_of[d] for d in order)
+    new_subscripts = ",".join(new_in_specs) + "->" + out_prefix + out_part
+
+    result = jax.numpy.einsum(new_subscripts, *bound, **kwargs)
+
+    # unbind: leading axes correspond to `order`, reindex them back to named
+    reindexed = jax_getitem(
+        result,
+        tuple(d() for d in order) + tuple(slice(None) for _ in range(len(out_part))),
+    )
+    return reindexed
+
+
 @_register_jax_op
 def jax_getitem(x: jax.Array, key: tuple[IndexElement, ...]) -> jax.Array:
     """Operation for indexing an array. Unlike the standard __getitem__ method,
@@ -215,6 +315,8 @@ def bind_dims[T, A, B](
     >>> bind_dims(t, b, a).shape
     (3, 2)
     """
+    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)
 
@@ -467,31 +467,66 @@ def _bind_dims_array(t: jax.Array, *args: Operation[[], jax.Array]) -> jax.Array
     array = t.args[0]
     dims = t.args[1]
     assert isinstance(dims, Sequence)
+    ndim = len(array.shape)
 
     # ensure that the order is a subset of the named dimensions
     order_set = set(args)
     if not order_set <= set(a.op for a in dims if isinstance(a, Term)):
         raise NotHandled
 
-    # permute the inner array so that the leading dimensions are in the order
-    # specified and the trailing dimensions are the remaining named dimensions
-    # (or slices)
-    reindex_dims = [
-        i
-        for i, o in enumerate(dims)
-        if not isinstance(o, Term) or o.op not in order_set
-    ]
-    dim_ops = [a.op if isinstance(a, Term) else None for a in dims]
-    perm = (
-        [dim_ops.index(o) for o in args]
-        + reindex_dims
-        + list(range(len(dims), len(array.shape)))
+    def axis_op(ax: int) -> Operation | None:
+        """The named op of a bare index term at axis ``ax``, else ``None``."""
+        if ax < len(dims):
+            d = dims[ax]
+            if isinstance(d, Term) and not d.args and not d.kwargs:
+                return d.op
+        return None
+
+    # Assign an einsum id to every axis of ``array``. Axes that share a named op
+    # get the *same* id — a repeated op (e.g. ``arr[i(), i()]``) ties its axes
+    # together, which einsum reads as a diagonal. Every other axis (slices, ints,
+    # fancy indices, compound terms, and trailing positional axes) gets a unique
+    # id, so einsum simply carries it through to be reindexed below.
+    op_ids: dict[Operation, int] = {}
+    in_ids: list[int] = []
+    next_id = 0
+    for ax in range(ndim):
+        op = axis_op(ax)
+        if op is not None:
+            if op not in op_ids:
+                op_ids[op] = next_id
+                next_id += 1
+            in_ids.append(op_ids[op])
+        else:
+            in_ids.append(next_id)
+            next_id += 1
+
+    # Output order: bound args that actually appear (in the requested order,
+    # deduplicated by the diagonal merge), then every remaining axis in
+    # first-appearance order. einsum does the permutation and the diagonals; with
+    # all distinct ids retained in the output it performs no reduction.
+    present_arg_ids = [op_ids[o] for o in args if o in op_ids]
+    seen = set(present_arg_ids)
+    rest_ids: list[int] = []
+    for i in in_ids:
+        if i not in seen:
+            seen.add(i)
+            rest_ids.append(i)
+
+    array = jnp.einsum(array, in_ids, present_arg_ids + rest_ids)
+
+    # Re-apply the original index for each carried axis and re-name unbound op
+    # axes. Trailing positional axes (first appearance beyond ``dims``) are left
+    # for jax_getitem to carry implicitly.
+    first_pos: dict[int, int] = {}
+    for ax, i in enumerate(in_ids):
+        first_pos.setdefault(i, ax)
+
+    index_expr = (slice(None),) * len(present_arg_ids) + tuple(
+        dims[first_pos[i]] if first_pos[i] < len(dims) else slice(None)
+        for i in rest_ids
     )
-    array = jnp.transpose(array, perm)
-    reindexed = jax_getitem(
-        array, (slice(None),) * len(args) + tuple(dims[i] for i in reindex_dims)
-    )
-    return reindexed
+    return jax_getitem(array, index_expr)
 
 
 @unbind_dims.register(jax.Array)  # type: ignore