effectful 0.1.0__py3-none-any.whl → 0.2.1__py3-none-any.whl

effectful/ops/types.py

@@ -2,20 +2,21 @@ from __future__ import annotations
 
 import abc
 import collections.abc
+import functools
 import inspect
 import typing
-from typing import Any, Callable, Generic, Mapping, Sequence, Type, TypeVar, Union
+from collections.abc import Callable, Mapping, Sequence
+from typing import Any, _ProtocolMeta, overload, runtime_checkable
 
-from typing_extensions import ParamSpec
 
-P = ParamSpec("P")
-Q = ParamSpec("Q")
-S = TypeVar("S")
-T = TypeVar("T")
-V = TypeVar("V")
+class NotHandled(Exception):
+    """Raised by an operation when the operation should remain unhandled."""
 
+    pass
 
-class Operation(abc.ABC, Generic[Q, V]):
+
+@functools.total_ordering
+class Operation[**Q, V](abc.ABC):
     """An abstract class representing an effect that can be implemented by an effect handler.
 
     .. note::
@@ -36,7 +37,11 @@ class Operation(abc.ABC, Generic[Q, V]):
         raise NotImplementedError
 
     @abc.abstractmethod
-    def __default_rule__(self, *args: Q.args, **kwargs: Q.kwargs) -> "Expr[V]":
+    def __lt__(self, other):
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def __default_rule__(self, *args: Q.args, **kwargs: Q.kwargs) -> Expr[V]:
         """The default rule is used when the operation is not handled.
 
         If no default rule is supplied, the free rule is used instead.
@@ -44,15 +49,12 @@ class Operation(abc.ABC, Generic[Q, V]):
         raise NotImplementedError
 
     @abc.abstractmethod
-    def __type_rule__(self, *args: Q.args, **kwargs: Q.kwargs) -> Type[V]:
+    def __type_rule__(self, *args: Q.args, **kwargs: Q.kwargs) -> type[V]:
         """Returns the type of the operation applied to arguments."""
         raise NotImplementedError
 
     @abc.abstractmethod
-    def __fvs_rule__(self, *args: Q.args, **kwargs: Q.kwargs) -> tuple[
-        tuple[collections.abc.Set["Operation"], ...],
-        dict[str, collections.abc.Set["Operation"]],
-    ]:
+    def __fvs_rule__(self, *args: Q.args, **kwargs: Q.kwargs) -> inspect.BoundArguments:
         """
         Returns the sets of variables that appear free in each argument and keyword argument
         but not in the result of the operation, i.e. the variables bound by the operation.
@@ -63,19 +65,17 @@ class Operation(abc.ABC, Generic[Q, V]):
         """
         raise NotImplementedError
 
-    @abc.abstractmethod
-    def __repr_rule__(self, *args: Q.args, **kwargs: Q.kwargs) -> str:
-        raise NotImplementedError
-
     @typing.final
     def __call__(self, *args: Q.args, **kwargs: Q.kwargs) -> V:
-        from effectful.internals.runtime import get_interpretation
         from effectful.ops.semantics import apply
 
-        return apply.__default_rule__(get_interpretation(), self, *args, **kwargs)  # type: ignore
+        return apply.__default_rule__(self, *args, **kwargs)  # type: ignore
 
+    def __repr__(self):
+        return f"{self.__class__.__name__}({self.__name__}, {self.__signature__})"
 
-class Term(abc.ABC, Generic[T]):
+
+class Term[T](abc.ABC):
     """A term in an effectful computation is a is a tree of :class:`Operation`
     applied to values.
 
@@ -91,33 +91,125 @@ class Term(abc.ABC, Generic[T]):
 
     @property
     @abc.abstractmethod
-    def args(self) -> Sequence["Expr[Any]"]:
+    def args(self) -> Sequence[Expr[Any]]:
         """Abstract property for the arguments."""
         raise NotImplementedError
 
     @property
     @abc.abstractmethod
-    def kwargs(self) -> Mapping[str, "Expr[Any]"]:
+    def kwargs(self) -> Mapping[str, Expr[Any]]:
         """Abstract property for the keyword arguments."""
         raise NotImplementedError
 
     def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.op!r}, {self.args!r}, {self.kwargs!r})"
+
+    def __str__(self) -> str:
         from effectful.internals.runtime import interpreter
         from effectful.ops.semantics import apply, evaluate
 
-        with interpreter({apply: lambda _, op, *a, **k: op.__repr_rule__(*a, **k)}):
-            return evaluate(self)  # type: ignore
+        fresh: dict[str, dict[Operation, int]] = collections.defaultdict(dict)
+
+        def op_str(op):
+            """Return a unique (in this term) name for the operation."""
+            name = op.__name__
+            if name not in fresh:
+                fresh[name] = {op: 0}
+            if op not in fresh[name]:
+                fresh[name][op] = len(fresh[name])
+
+            n = fresh[name][op]
+            if n == 0:
+                return name
+            return f"{name}!{n}"
+
+        def term_str(term):
+            if isinstance(term, Operation):
+                return op_str(term)
+            elif isinstance(term, list):
+                return "[" + ", ".join(map(term_str, term)) + "]"
+            elif isinstance(term, tuple):
+                return "(" + ", ".join(map(term_str, term)) + ")"
+            elif isinstance(term, dict):
+                return (
+                    "{"
+                    + ", ".join(
+                        f"{term_str(k)}:{term_str(v)}" for (k, v) in term.items()
+                    )
+                    + "}"
+                )
+            return str(term)
+
+        def _apply(op, *args, **kwargs) -> str:
+            args_str = ", ".join(map(term_str, args)) if args else ""
+            kwargs_str = (
+                ", ".join(f"{k}={term_str(v)}" for k, v in kwargs.items())
+                if kwargs
+                else ""
+            )
+
+            ret = f"{op_str(op)}({args_str}"
+            if kwargs:
+                ret += f"{', ' if args else ''}"
+            ret += f"{kwargs_str})"
+            return ret
+
+        with interpreter({apply: _apply}):
+            return typing.cast(str, evaluate(self))
 
 
 #: An expression is either a value or a term.
-Expr = Union[T, Term[T]]
+type Expr[T] = T | Term[T]
 
-#: An interpretation is a mapping from operations to their implementations.
-Interpretation = Mapping[Operation[..., T], Callable[..., V]]
 
+class _InterpretationMeta(_ProtocolMeta):
+    def __instancecheck__(cls, instance):
+        return isinstance(instance, collections.abc.Mapping) and all(
+            isinstance(k, Operation) and callable(v) for k, v in instance.items()
+        )
 
-class Annotation(abc.ABC):
 
+@runtime_checkable
+class Interpretation[T, V](typing.Protocol, metaclass=_InterpretationMeta):
+    """An interpretation is a mapping from operations to their implementations."""
+
+    def keys(self):
+        raise NotImplementedError
+
+    def values(self):
+        raise NotImplementedError
+
+    def items(self):
+        raise NotImplementedError
+
+    @overload
+    def get(self, key: Operation[..., T], /) -> Callable[..., V] | None:
+        raise NotImplementedError
+
+    @overload
+    def get(
+        self, key: Operation[..., T], default: Callable[..., V], /
+    ) -> Callable[..., V]:
+        raise NotImplementedError
+
+    @overload
+    def get[S](self, key: Operation[..., T], default: S, /) -> Callable[..., V] | S:
+        raise NotImplementedError
+
+    def __getitem__(self, key: Operation[..., T]) -> Callable[..., V]:
+        raise NotImplementedError
+
+    def __contains__(self, key: Operation[..., T]) -> bool:
+        raise NotImplementedError
+
+    def __iter__(self):
+        raise NotImplementedError
+
+    def __len__(self) -> int:
+        raise NotImplementedError
+
+
+class Annotation(abc.ABC):
     @classmethod
     @abc.abstractmethod
     def infer_annotations(cls, sig: inspect.Signature) -> inspect.Signature:

{effectful-0.1.0.dist-info → effectful-0.2.1.dist-info}/METADATA

@@ -1,72 +1,59 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: effectful
-Version: 0.1.0
+Version: 0.2.1
 Summary: Metaprogramming infrastructure
-Home-page: https://www.basis.ai/
 Author: Basis
-License: Apache 2.0
+License-Expression: Apache-2.0
+Project-URL: Homepage, https://www.basis.ai/
 Project-URL: Source, https://github.com/BasisResearch/effectful
-Keywords: machine learning statistics probabilistic programming bayesian modeling pytorch
+Project-URL: Bug Tracker, https://github.com/BasisResearch/effectful/issues
+Keywords: machine learning,statistics,probabilistic programming,bayesian modeling,pytorch
 Classifier: Intended Audience :: Developers
 Classifier: Intended Audience :: Education
 Classifier: Intended Audience :: Science/Research
-Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: POSIX :: Linux
 Classifier: Operating System :: MacOS :: MacOS X
-Classifier: Programming Language :: Python :: 3.10
-Classifier: Programming Language :: Python :: 3.11
-Requires-Python: >=3.10
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Python: >=3.12
+Description-Content-Type: text/x-rst
 License-File: LICENSE.md
-Requires-Dist: typing_extensions
-Requires-Dist: dm-tree
 Provides-Extra: torch
 Requires-Dist: torch; extra == "torch"
+Requires-Dist: dm-tree; extra == "torch"
 Provides-Extra: pyro
-Requires-Dist: torch; extra == "pyro"
-Requires-Dist: pyro-ppl; extra == "pyro"
-Provides-Extra: dev
-Requires-Dist: torch; extra == "dev"
-Requires-Dist: pyro-ppl; extra == "dev"
-Requires-Dist: torch; extra == "dev"
-Requires-Dist: setuptools; extra == "dev"
-Requires-Dist: sphinx; extra == "dev"
-Requires-Dist: sphinxcontrib-bibtex; extra == "dev"
-Requires-Dist: sphinx_rtd_theme; extra == "dev"
-Requires-Dist: myst-parser; extra == "dev"
-Requires-Dist: nbsphinx; extra == "dev"
-Requires-Dist: pytest; extra == "dev"
-Requires-Dist: pytest-cov; extra == "dev"
-Requires-Dist: pytest-xdist; extra == "dev"
-Requires-Dist: pytest-benchmark; extra == "dev"
-Requires-Dist: mypy; extra == "dev"
-Requires-Dist: black; extra == "dev"
-Requires-Dist: flake8; extra == "dev"
-Requires-Dist: isort; extra == "dev"
-Requires-Dist: nbval; extra == "dev"
-Requires-Dist: nbqa; extra == "dev"
+Requires-Dist: pyro-ppl>=1.9.1; extra == "pyro"
+Requires-Dist: dm-tree; extra == "pyro"
+Provides-Extra: jax
+Requires-Dist: jax; extra == "jax"
+Requires-Dist: dm-tree; extra == "jax"
+Provides-Extra: numpyro
+Requires-Dist: numpyro>=0.19; extra == "numpyro"
+Requires-Dist: dm-tree; extra == "numpyro"
 Provides-Extra: docs
-Requires-Dist: setuptools; extra == "docs"
+Requires-Dist: effectful[jax,numpyro,pyro,torch]; extra == "docs"
 Requires-Dist: sphinx; extra == "docs"
 Requires-Dist: sphinxcontrib-bibtex; extra == "docs"
 Requires-Dist: sphinx_rtd_theme; extra == "docs"
 Requires-Dist: myst-parser; extra == "docs"
 Requires-Dist: nbsphinx; extra == "docs"
-Dynamic: author
-Dynamic: classifier
-Dynamic: description
-Dynamic: home-page
-Dynamic: keywords
-Dynamic: license
-Dynamic: project-url
-Dynamic: provides-extra
-Dynamic: requires-dist
-Dynamic: requires-python
-Dynamic: summary
-
+Requires-Dist: sphinx_autodoc_typehints; extra == "docs"
+Requires-Dist: pypandoc_binary; extra == "docs"
+Provides-Extra: test
+Requires-Dist: effectful[docs,jax,numpyro,pyro,torch]; extra == "test"
+Requires-Dist: pytest; extra == "test"
+Requires-Dist: pytest-cov; extra == "test"
+Requires-Dist: pytest-xdist; extra == "test"
+Requires-Dist: pytest-benchmark; extra == "test"
+Requires-Dist: mypy; extra == "test"
+Requires-Dist: ruff; extra == "test"
+Requires-Dist: nbval; extra == "test"
+Requires-Dist: nbqa; extra == "test"
+Dynamic: license-file
 
 .. index-inclusion-marker
 
-Effectful 
+Effectful
 =========
 
 Effectful is an algebraic effect system for Python, intended for use in the
@@ -89,9 +76,12 @@ Install From Source
 Install With Optional PyTorch/Pyro Support
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-``effectful`` has optional support for `PyTorch <https://pytorch.org/>`_ (tensors
-with named dimensions) and `Pyro <https://pyro.ai/>`_ (wrappers for Pyro
-effects).
+``effectful`` has optional support for:
+
+- `PyTorch <https://pytorch.org/>`_ (tensors with named dimensions)
+- `Pyro <https://pyro.ai/>`_ (wrappers for Pyro effects)
+- `Jax <https://docs.jax.dev/en/latest/index.html>`_ (tensors with named dimensions)
+- `Numpyro <https://num.pyro.ai>`_ (operations for Numpyro distributions)
 
 To enable PyTorch support:
 
@@ -105,6 +95,18 @@ Pyro support (which includes PyTorch support):
 
    pip install effectful[pyro]
 
+Jax support:
+
+.. code:: sh
+
+   pip install effectful[jax]
+
+Numpyro support (which includes Jax support):
+
+.. code:: sh
+
+   pip install effectful[numpyro]
+
 Getting Started
 ---------------
 
@@ -115,11 +117,12 @@ Here's an example demonstrating how ``effectful`` can be used to implement a sim
    import functools
 
    from effectful.ops.types import Term
-   from effectful.ops.syntax import defop
+   from effectful.ops.syntax import defdata, defop
    from effectful.ops.semantics import handler, evaluate, coproduct, fwd
-   from effectful.handlers.numbers import add
 
-   def beta_add(x: int, y: int) -> int:        
+   add = defdata.dispatch(int).__add__
+
+   def beta_add(x: int, y: int) -> int:
        match x, y:
            case int(), int():
                return x + y
@@ -129,14 +132,14 @@ Here's an example demonstrating how ``effectful`` can be used to implement a sim
    def commute_add(x: int, y: int) -> int:
        match x, y:
            case Term(), int():
-               return y + x  
+               return y + x
            case _:
                return fwd()
 
    def assoc_add(x: int, y: int) -> int:
        match x, y:
            case _, Term(op, (a, b)) if op == add:
-               return (x + a) + b 
+               return (x + a) + b
            case _:
                return fwd()
 
@@ -168,7 +171,7 @@ We can make the evaluation strategy smarter by taking advantage of the commutati
    >>> with handler(eager_mixed):
    >>>     print(evaluate(e))
    add(8, add(x(), y()))
-   
+
 Learn More
 ----------
 

effectful-0.2.1.dist-info/RECORD

@@ -0,0 +1,26 @@
+effectful/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+effectful/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+effectful/handlers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+effectful/handlers/indexed.py,sha256=ZY-8w32a1PKGVScwXjbgByI3wRHvfxuuXJVwLlp0rgw,12622
+effectful/handlers/numpyro.py,sha256=RWoBNpHLr5KdotN5Vu118jY0kn_p6NaejcnJgZJLehw,19457
+effectful/handlers/pyro.py,sha256=qVl1wson02pyV8YHGf93KDnYEp5pGmhKEwji95OYBl8,26486
+effectful/handlers/torch.py,sha256=NNM7mxqZskEBCjsl25kHI95WlXG9aeD7FaSkXkoLZ_I,24330
+effectful/handlers/jax/__init__.py,sha256=O-BygB2HCkDftP1B98mQnsrt7sOdZpvC2YoHXd61tmI,494
+effectful/handlers/jax/_handlers.py,sha256=95fGOZoTQfYehs0ip_zR91TKwnJB-beK6wr_SH73Bg8,9252
+effectful/handlers/jax/_terms.py,sha256=k_nLe5jvj76ZHVr4LWWFbQT4lxhGUVTEWQB6W38_M8I,16341
+effectful/handlers/jax/numpy/__init__.py,sha256=Kmvya0QI-GA56pPf1as-wYOuZFngOBLtsawa35vPKhg,516
+effectful/handlers/jax/numpy/linalg.py,sha256=9DiaYYG4SztmO-VkmMH3dVvULtMK-zEgbV9oNQFkFo8,350
+effectful/handlers/jax/scipy/special.py,sha256=yTIECFtQVPgraonrPlyenjvcnEYchZwIZC-5CSkF-lA,299
+effectful/internals/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+effectful/internals/runtime.py,sha256=aLWol7sR1yHekn7zNz1evHKHARjiT1tnkmByLHPHBGc,1811
+effectful/internals/tensor_utils.py,sha256=3QCSUqdxCXod3dsY3oRMcg36Rqr8pVX-ktEyCEkeODo,1173
+effectful/internals/unification.py,sha256=MD0beZ29by9hbJaowafBLIZuKPF5MxHXezuGpSG_wCA,30254
+effectful/ops/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+effectful/ops/semantics.py,sha256=u6i7RTpiv5fCnOdh_pOL-5CSsO4QtW-SjfGEQiizeWk,11922
+effectful/ops/syntax.py,sha256=PE9iuLIzwqFkkk99bBZWk93X4sqSUhtRPePiTQRnAdk,55825
+effectful/ops/types.py,sha256=W1gZJaBnX7_nFpWrG3vfCBQPSun3Gc9PqT61ls8B3EA,6599
+effectful-0.2.1.dist-info/licenses/LICENSE.md,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+effectful-0.2.1.dist-info/METADATA,sha256=AgFY63eWiiFpOlvxgWa-3GOV9eSy__N0gF6ZIs5ljdc,5300
+effectful-0.2.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+effectful-0.2.1.dist-info/top_level.txt,sha256=gtuJfrE2nXil_lZLCnqWF2KAbOnJs9ILNvK8WnkRzbs,10
+effectful-0.2.1.dist-info/RECORD,,

{effectful-0.1.0.dist-info → effectful-0.2.1.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.8.0)
+Generator: setuptools (80.9.0)
 Root-Is-Purelib: true
 Tag: py3-none-any