Add collection constructor operations and rewrite analyses to remove side effects

effectful/handlers/jax/_handlers.py

@@ -1,4 +1,5 @@
 import functools
+import itertools
 import typing
 from collections.abc import Callable, Mapping, Sequence
 from types import EllipsisType
@@ -10,10 +11,12 @@
 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.syntax import (
+    ConstructorOperation,
+    PureInterpretation,
     Scoped,
+    _BaseTerm,
     _CustomSingleDispatchCallable,
     defdata,
     deffn,
@@ -39,7 +42,71 @@ def is_eager_array(x):
     )
 
 
-def sizesof(value) -> Mapping[Operation[[], jax.Array], int]:
+@functools.cache
+def _sizesof_intp() -> tuple[PureInterpretation, Operation]:
+    """Construct the singleton interpretation used by ``sizesof``."""
+    from effectful.internals.product_n import argsof, productN
+
+    _sizes = defop(object, name="sizes")
+    _getitem_term = defop(object, name="getitem_args")
+
+    def _retain(op, *args, **kwargs):
+        # Non-getitem subterms are opaque to this analysis. Keeping their
+        # arguments would retain the entire input term unnecessarily.
+        return _BaseTerm(op)
+
+    def _retain_getitem(*args, **kwargs):
+        return defdata(jax_getitem, *args, **kwargs)
+
+    def _merge(s1, s2):
+        s3 = s1.copy()
+        for k, v in s2.items():
+            if k in s3 and s3[k] != v:
+                raise ValueError(
+                    f"Named index {k} used in incompatible dimensions of size {s3[k]} and {v}"
+                )
+            s3[k] = v
+        return s3
+
+    def _apply_sizes(op, *args, **kwargs):
+        analyses = (x for x in (*args, *kwargs.values()) if isinstance(x, dict))
+        return functools.reduce(_merge, analyses, {})
+
+    def _getitem(arr, index):
+        # Inspect this getitem's arguments in the term projection without
+        # forcing that projection to retain the getitem result.
+        term_args, _ = argsof(_getitem_term)
+        term_arr, term_index = term_args
+
+        arg_sizes = (x for x in (arr, index) if isinstance(x, dict))
+        if not is_eager_array(term_arr):
+            return functools.reduce(_merge, arg_sizes, {})
+
+        sizes = (
+            {k.op: term_arr.shape[i]}
+            for i, k in enumerate(term_index)
+            if isinstance(k, Term) and len(k.args) == 0 and len(k.kwargs) == 0
+        )
+        return functools.reduce(_merge, itertools.chain(arg_sizes, sizes), {})
+
+    return (
+        PureInterpretation(
+            productN(
+                {
+                    _sizes: {apply: _apply_sizes, jax_getitem: _getitem},
+                    _getitem_term: {
+                        apply: _retain,
+                        jax_getitem: _retain_getitem,
+                        ConstructorOperation.__apply__: apply.__default_rule__,
+                    },
+                }
+            )
+        ),
+        _sizes,
+    )
+
+
+def sizesof(term: Expr) -> Mapping[Operation[[], jax.Array], int]:
     """Return the sizes of named dimensions in an array expression.
 
     Sizes are inferred from the array shape.
@@ -53,30 +120,14 @@ def sizesof(value) -> Mapping[Operation[[], jax.Array], int]:
     >>> sizes = sizesof(jax_getitem(jnp.ones((2, 3)), [a(), b()]))
     >>> assert sizes[a] == 2 and sizes[b] == 3
     """
-    sizes: dict[Operation[[], jax.Array], int] = {}
-
-    def update_sizes(sizes, op, size):
-        old_size = sizes.get(op)
-        if old_size is not None and size != old_size:
-            raise ValueError(
-                f"Named index {op} used in incompatible dimensions of size {old_size} and {size}"
-            )
-        sizes[op] = size
-
-    def _getitem_sizeof(x: jax.Array, key: tuple[Expr[IndexElement], ...]):
-        if is_eager_array(x):
-            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])
-        return defdata(jax_getitem, x, key)
-
-    def _apply(op, *args, **kwargs):
-        return defdata(op, *args, **kwargs)
-
-    with interpreter({jax_getitem: _getitem_sizeof, apply: _apply}):
-        evaluate(value)
-
-    return sizes
+    from effectful.internals.product_n import _unpack
+
+    intp, prompt = _sizesof_intp()
+    result = evaluate(term, intp=intp)
+    fvs = _unpack(result, prompt)
+    if not isinstance(fvs, dict):
+        return {}
+    return fvs
 
 
 def _partial_eval(t: Expr[jax.Array]) -> Expr[jax.Array]:

effectful/handlers/jax/numpy/__init__.py

@@ -1,3 +1,4 @@
+import types
 from typing import TYPE_CHECKING
 
 import jax.numpy
@@ -7,17 +8,28 @@
 _no_overload = ["array", "asarray"]
 
 for name, op in jax.numpy.__dict__.items():
-    if not callable(op):
+    if isinstance(op, types.ModuleType):
         continue
 
-    jax_op = (
-        _register_jax_op_no_partial_eval(op)
-        if name in _no_overload
-        else _register_jax_op(op)
-    )
-    globals()[name] = jax_op
+    # copy constants
+    if isinstance(op, float | types.NoneType):
+        globals()[name] = op
 
-pi = jax.numpy.pi
+    if callable(op):
+        if name == "__getattr__":
+            continue
+
+        elif name in _no_overload:
+            globals()[name] = _register_jax_op_no_partial_eval(op)
+
+        else:
+            globals()[name] = _register_jax_op(op)
+        jax_op = (
+            _register_jax_op_no_partial_eval(op)
+            if name in _no_overload
+            else _register_jax_op(op)
+        )
+        globals()[name] = jax_op
 
 # Tell mypy about our wrapped functions.
 if TYPE_CHECKING:

effectful/handlers/pyro.py

@@ -28,7 +28,7 @@
 )
 from effectful.internals.runtime import interpreter
 from effectful.ops.semantics import apply, evaluate, handler, typeof
-from effectful.ops.syntax import defdata, defop
+from effectful.ops.syntax import ConstructorOperation, defdata, defop
 from effectful.ops.types import NotHandled, Operation, Term
 
 
@@ -368,7 +368,9 @@ def _to_named(a):
             return a
 
     # Convert to a term in a context that does not evaluate distribution constructors.
-    with handler({apply: defdata}):
+    with handler(
+        {apply: defdata, ConstructorOperation.__apply__: apply.__default_rule__}
+    ):
         d = typing.cast(TorchDistribution, evaluate(value))
 
     if not (isinstance(d, Term) and typeof(d) is TorchDistribution):
@@ -403,7 +405,9 @@ def _to_positional(a, indices):
         else:
             return a
 
-    with handler({apply: defdata}):
+    with handler(
+        {apply: defdata, ConstructorOperation.__apply__: apply.__default_rule__}
+    ):
         d = typing.cast(TorchDistribution, evaluate(value))
 
     if not (isinstance(d, Term) and typeof(d) is TorchDistribution):

effectful/handlers/torch.py

@@ -1,4 +1,5 @@
 import functools
+import itertools
 import typing
 from collections.abc import Callable, Mapping, Sequence
 from types import EllipsisType
@@ -11,10 +12,17 @@
 
 import torch.utils._pytree as pytree
 
-from effectful.internals.runtime import interpreter
 from effectful.internals.tensor_utils import _desugar_tensor_index
 from effectful.ops.semantics import apply, evaluate, fvsof, handler, typeof
-from effectful.ops.syntax import Scoped, defdata, defop, syntactic_eq
+from effectful.ops.syntax import (
+    ConstructorOperation,
+    PureInterpretation,
+    Scoped,
+    _BaseTerm,
+    defdata,
+    defop,
+    syntactic_eq,
+)
 from effectful.ops.types import Expr, NotHandled, Operation, Term
 
 # + An element of a tensor index expression.
@@ -34,6 +42,74 @@ def _getitem_ellipsis_and_none(
     return torch.reshape(x, new_shape), new_key
 
 
+@functools.cache
+def _sizesof_intp() -> tuple[PureInterpretation, Operation]:
+    """Construct the singleton interpretation used by ``sizesof``."""
+    from effectful.internals.product_n import argsof, productN
+
+    sizes = defop(object, name="sizes")
+    getitem_term = defop(object, name="getitem_args")
+
+    def _retain(op, *args, **kwargs):
+        # Non-getitem subterms are opaque to this analysis. Keeping their
+        # arguments would retain the entire input term unnecessarily.
+        return _BaseTerm(op)
+
+    def _retain_getitem(*args, **kwargs):
+        return defdata(torch_getitem, *args, **kwargs)
+
+    def _merge(s1, s2):
+        result = s1.copy()
+        for k, v in s2.items():
+            if k in result and result[k] != v:
+                raise ValueError(
+                    f"Named index {k} used in incompatible dimensions of size {result[k]} and {v}"
+                )
+            result[k] = v
+        return result
+
+    def _apply_sizes(op, *args, **kwargs):
+        analyses = (x for x in (*args, *kwargs.values()) if isinstance(x, dict))
+        return functools.reduce(_merge, analyses, {})
+
+    def _getitem(x, key):
+        # Inspect this getitem's arguments in the term projection without
+        # forcing that projection to retain the getitem result.
+        term_args, _ = argsof(getitem_term)
+        term_x, term_key = term_args
+
+        arg_sizes = (value for value in (x, key) if isinstance(value, dict))
+        if not isinstance(term_x, torch.Tensor):
+            return functools.reduce(_merge, arg_sizes, {})
+
+        shape, desugared_key = _desugar_tensor_index(term_x.shape, term_key)
+        index_sizes = (
+            {k.op: shape[i]}
+            for i, k in enumerate(desugared_key)
+            if isinstance(k, Term)
+            and not k.args
+            and not k.kwargs
+            and issubclass(typeof(k), torch.Tensor)
+        )
+        return functools.reduce(_merge, itertools.chain(arg_sizes, index_sizes), {})
+
+    return (
+        PureInterpretation(
+            productN(
+                {
+                    sizes: {apply: _apply_sizes, torch_getitem: _getitem},
+                    getitem_term: {
+                        apply: _retain,
+                        torch_getitem: _retain_getitem,
+                        ConstructorOperation.__apply__: apply.__default_rule__,
+                    },
+                }
+            )
+        ),
+        sizes,
+    )
+
+
 def sizesof(value) -> Mapping[Operation[[], torch.Tensor], int]:
     """Return the sizes of named dimensions in a tensor expression.
 
@@ -48,35 +124,13 @@ def sizesof(value) -> Mapping[Operation[[], torch.Tensor], int]:
     >>> sizes = sizesof(torch.ones(2, 3)[a(), b()])
     >>> assert sizes[a] == 2 and sizes[b] == 3
     """
-    sizes: dict[Operation[[], torch.Tensor], int] = {}
-
-    def _torch_getitem_sizeof(
-        x: Expr[torch.Tensor], key: tuple[Expr[IndexElement], ...]
-    ) -> Expr[torch.Tensor]:
-        if isinstance(x, torch.Tensor):
-            shape, key_ = _desugar_tensor_index(x.shape, key)
-
-            for i, k in enumerate(key_):
-                if (
-                    isinstance(k, Term)
-                    and len(k.args) == 0
-                    and len(k.kwargs) == 0
-                    and issubclass(typeof(k), torch.Tensor)
-                ):
-                    if k.op in sizes and sizes[k.op] != shape[i]:
-                        raise ValueError(
-                            f"Named index {k.op} used in incompatible dimensions of size {sizes[k.op]} and {shape[i]}"
-                        )
-                    sizes[k.op] = shape[i]
-
-        return defdata(torch_getitem, x, key)
-
-    def _apply(op, *args, **kwargs):
-        return defdata(op, *args, **kwargs)
-
-    with interpreter({torch_getitem: _torch_getitem_sizeof, apply: _apply}):
-        evaluate(value)
+    from effectful.internals.product_n import _unpack
 
+    intp, prompt = _sizesof_intp()
+    result = evaluate(value, intp=intp)
+    sizes = _unpack(result, prompt)
+    if not isinstance(sizes, dict):
+        return {}
     return sizes