Staging branch for weighted

docs/source/introduction.ipynb

@@ -21,15 +21,15 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "id": "5278fd54",
    "metadata": {},
    "outputs": [],
    "source": [
     "import functools\n",
     "\n",
     "from effectful.ops.semantics import coproduct, fwd, handler\n",
-    "from effectful.ops.syntax import defdata, defop\n",
+    "from effectful.ops.syntax import defdata, defop, syntactic_eq\n",
     "from effectful.ops.types import NotHandled, Operation, Term\n",
     "\n",
     "\n",
@@ -65,7 +65,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 2,
    "id": "3c575e02",
    "metadata": {
     "lines_to_next_cell": 2
@@ -78,7 +78,7 @@
     "\n",
     "\n",
     "def test_adding_two_numbers():\n",
-    "    assert add(1, 3) == add(2, 2)\n",
+    "    assert syntactic_eq(add(1, 3), add(2, 2))\n",
     "\n",
     "\n",
     "assert isinstance(add, Operation)\n",
@@ -104,7 +104,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 3,
    "id": "6e293b33",
    "metadata": {},
    "outputs": [],
@@ -183,7 +183,7 @@
     "\n",
     "\n",
     "def test_adding_a_variable():\n",
-    "    assert add(x, 2) == add(x, 2)\n",
+    "    assert syntactic_eq(add(x, 2), add(x, 2))\n",
     "\n",
     "\n",
     "fails(test_adding_a_variable)\n",
@@ -374,7 +374,7 @@
     "\n",
     "\n",
     "def test_mixed_together():\n",
-    "    assert add(add(x, 2), add(2, y)) == add(x, add(add(3, y), 1))\n",
+    "    assert syntactic_eq(add(add(x, 2), add(2, y)), add(x, add(add(3, y), 1)))\n",
     "\n",
     "\n",
     "fails(test_mixed_together)\n",
@@ -512,7 +512,7 @@
    "notebook_metadata_filter": "-all"
   },
   "kernelspec": {
-   "display_name": "base",
+   "display_name": "effectful (3.12.9.final.0)",
    "language": "python",
    "name": "python3"
   },
@@ -526,7 +526,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.13"
+   "version": "3.12.9"
   }
  },
  "nbformat": 4,

effectful/handlers/jax/_handlers.py

@@ -1,16 +1,21 @@
 import functools
+import itertools
 import typing
-from collections.abc import Callable, Mapping
+from collections.abc import Callable, Iterable, Mapping
 from typing import Annotated
 
+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")
 
 from effectful.internals.tensor_utils import IndexElement, _BaseSizesofIntp, _sizesof
-from effectful.ops.semantics import fvsof, typeof
+from effectful.ops.semantics import fvsof, fwd, typeof
 from effectful.ops.syntax import (
     Scoped,
     _CustomSingleDispatchCallable,
@@ -19,10 +24,14 @@
     defop,
     implements,
     syntactic_eq,
+    syntactic_hash,
 )
 from effectful.ops.types import Expr, NotHandled, Operation, Term
 
 
+class JaxOperation[**P, T](Operation[P, T]): ...
+
+
 def is_eager_array(x):
     return isinstance(x, jax.Array) or (
         isinstance(x, Term)
@@ -43,6 +52,13 @@ def _partial_eval(t: Expr[jax.Array]) -> Expr[jax.Array]:
     if not sized_fvs:
         return t
 
+    # if any dimension is zero sized, the result is empty
+    if any(size == 0 for size in sized_fvs.values()):
+        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):
         return not isinstance(t, Term) or t.op in sized_fvs or is_eager_array(t)
 
@@ -82,7 +98,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)
@@ -91,10 +107,17 @@ def _jax_op(*args, **kwargs) -> jax.Array:
             _jax_op is jax_getitem
             and not isinstance(args[0], Term)
             and sized_fvs
+            and not isinstance(args[1], Term)
             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 _jax_op is jax_getitem and not isinstance(args[1], Term) 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,
@@ -132,6 +155,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, Term) and not is_eager_array(arr) 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], ...]) -> jax.Array:
     """Operation for indexing an array. Unlike the standard __getitem__ method,
@@ -193,11 +303,20 @@ 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)  # type: ignore[attr-defined]
+def _(*args, **kwargs):
+    raise NotHandled
+
+
 @defop
 @_CustomSingleDispatchCallable
 def unbind_dims[T, A, B](
@@ -206,11 +325,18 @@ def unbind_dims[T, A, B](
     *names: Annotated[Operation[[], jax.Array], 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)  # type: ignore[attr-defined]
+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)
@@ -262,3 +388,20 @@ def _(x: jax.Array, other) -> bool:
         and x.shape == other.shape
         and bool((jnp.asarray(x) == jnp.asarray(other)).all())
     )
+
+
+@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), x.tobytes()))
+
+
+@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)))