effectful 0.0.1__tar.gz → 0.2.0__tar.gz

{effectful-0.0.1 → effectful-0.2.0}/PKG-INFO

@@ -1,63 +1,59 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: effectful
-Version: 0.0.1
+Version: 0.2.0
 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: 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: 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: nbval; extra == "dev"
-Requires-Dist: nbqa; extra == "dev"
-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: 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: 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"
+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
@@ -75,14 +71,17 @@ Install From Source
    git clone git@github.com:BasisResearch/effectful.git
    cd effectful
    git checkout master
-   pip install -e .[test]
+   pip install -e .[pyro]
 
 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:
 
@@ -96,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
 ---------------
 
@@ -103,18 +114,15 @@ Here's an example demonstrating how ``effectful`` can be used to implement a sim
 
 .. code:: python
 
-   import operator
    import functools
 
-   from effectful.handlers.operator import OPERATORS
    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
-   import effectful.handlers.operator
 
-   add = OPERATORS[operator.add]
+   add = defdata.dispatch(int).__add__
 
-   def beta_add(x: int, y: int) -> int:        
+   def beta_add(x: int, y: int) -> int:
        match x, y:
            case int(), int():
                return x + y
@@ -124,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()
 
@@ -163,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.0.1 → effectful-0.2.0}/README.rst

@@ -1,6 +1,6 @@
 .. index-inclusion-marker
 
-Effectful 
+Effectful
 =========
 
 Effectful is an algebraic effect system for Python, intended for use in the
@@ -18,14 +18,17 @@ Install From Source
    git clone git@github.com:BasisResearch/effectful.git
    cd effectful
    git checkout master
-   pip install -e .[test]
+   pip install -e .[pyro]
 
 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:
 
@@ -39,6 +42,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
 ---------------
 
@@ -46,18 +61,15 @@ Here's an example demonstrating how ``effectful`` can be used to implement a sim
 
 .. code:: python
 
-   import operator
    import functools
 
-   from effectful.handlers.operator import OPERATORS
    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
-   import effectful.handlers.operator
 
-   add = OPERATORS[operator.add]
+   add = defdata.dispatch(int).__add__
 
-   def beta_add(x: int, y: int) -> int:        
+   def beta_add(x: int, y: int) -> int:
        match x, y:
            case int(), int():
                return x + y
@@ -67,14 +79,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()
 
@@ -106,7 +118,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.0.1 → effectful-0.2.0}/effectful/handlers/indexed.py

@@ -1,18 +1,16 @@
 import functools
 import operator
-from typing import Any, Dict, Iterable, Optional, Set, TypeVar, Union
+from collections.abc import Iterable
+from typing import Any
 
 import torch
 
-from effectful.handlers.torch import Indexable, sizesof
+from effectful.handlers.torch import sizesof
 from effectful.ops.syntax import deffn, defop
-from effectful.ops.types import Operation, Term
+from effectful.ops.types import Operation
 
-K = TypeVar("K")
-T = TypeVar("T")
 
-
-class IndexSet(Dict[str, Set[int]]):
+class IndexSet(dict[str, set[int]]):
     """
     :class:`IndexSet` s represent the support of an indexed value,
     for which free variables correspond to single interventions and indices
@@ -32,13 +30,13 @@ class IndexSet(Dict[str, Set[int]]):
     for which a value is defined::
 
         >>> IndexSet(x={0, 1}, y={2, 3})
-        IndexSet({x: {0, 1}, y: {2, 3}})
+        IndexSet({'x': {0, 1}, 'y': {2, 3}})
 
     :class:`IndexSet` 's constructor will automatically drop empty entries
     and attempt to convert input values to :class:`set` s::
 
         >>> IndexSet(x=[0, 0, 1], y=set(), z=2)
-        IndexSet({x: {0, 1}, z: {2}})
+        IndexSet({'x': {0, 1}, 'z': {2}})
 
     :class:`IndexSet` s are also hashable and can be used as keys in :class:`dict` s::
 
@@ -47,7 +45,7 @@ class IndexSet(Dict[str, Set[int]]):
         True
     """
 
-    def __init__(self, **mapping: Union[int, Iterable[int]]):
+    def __init__(self, **mapping: int | Iterable[int]):
         index_set = {}
         for k, vs in mapping.items():
             indexes = {vs} if isinstance(vs, int) else set(vs)
@@ -61,16 +59,6 @@ class IndexSet(Dict[str, Set[int]]):
     def __hash__(self):
         return hash(frozenset((k, frozenset(vs)) for k, vs in self.items()))
 
-    def _to_handler(self):
-        """Return an effectful handler that binds each index variable to a
-        tensor of its possible index values.
-
-        """
-        return {
-            name_to_sym(k): functools.partial(lambda v: v, torch.tensor(list(v)))
-            for k, v in self.items()
-        }
-
 
 def union(*indexsets: IndexSet) -> IndexSet:
     """
@@ -166,25 +154,17 @@ def indices_of(value: Any) -> IndexSet:
     :param kwargs: Additional keyword arguments used by specific implementations.
     :return: A :class:`IndexSet` containing the indices on which the value is supported.
     """
-    if isinstance(value, Term):
-        return IndexSet(
-            **{
-                k.__name__: set(range(v))  # type:ignore
-                for (k, v) in sizesof(value).items()
-            }
-        )
-    elif isinstance(value, torch.distributions.Distribution):
-        return indices_of(value.sample())
-
-    return IndexSet()
+    return IndexSet(
+        **{getattr(k, "__name__"): set(range(v)) for (k, v) in sizesof(value).items()}
+    )
 
 
-@functools.lru_cache(maxsize=None)
-def name_to_sym(name: str) -> Operation[[], int]:
-    return defop(int, name=name)
+@functools.cache
+def name_to_sym(name: str) -> Operation[[], torch.Tensor]:
+    return defop(torch.Tensor, name=name)
 
 
-def gather(value: torch.Tensor, indexset: IndexSet, **kwargs) -> torch.Tensor:
+def gather(value: torch.Tensor, indexset: IndexSet) -> torch.Tensor:
     """
     Selects entries from an indexed value at the indices in a :class:`IndexSet` .
     :func:`gather` is useful in conjunction with :class:`MultiWorldCounterfactual`
@@ -248,9 +228,7 @@ def gather(value: torch.Tensor, indexset: IndexSet, **kwargs) -> torch.Tensor:
     """
     indexset_vars = {name_to_sym(name): inds for name, inds in indexset.items()}
     binding = {
-        k: functools.partial(
-            lambda v: v, Indexable(torch.tensor(list(indexset_vars[k])))[k()]
-        )
+        k: functools.partial(lambda v: v, torch.tensor(list(indexset_vars[k]))[k()])
         for k in sizesof(value).keys()
         if k in indexset_vars
     }
@@ -259,14 +237,15 @@ def gather(value: torch.Tensor, indexset: IndexSet, **kwargs) -> torch.Tensor:
 
 
 def stack(
-    values: Union[tuple[torch.Tensor, ...], list[torch.Tensor]], name: str, **kwargs
+    values: tuple[torch.Tensor, ...] | list[torch.Tensor], name: str
 ) -> torch.Tensor:
     """Stack a sequence of indexed values, creating a new dimension. The new
     dimension is indexed by `dim`. The indexed values in the stack must have
     identical shapes.
 
     """
-    return Indexable(torch.stack(values))[name_to_sym(name)()]
+    values = torch.distributions.utils.broadcast_all(*values)
+    return torch.stack(values)[name_to_sym(name)()]
 
 
 def cond(fst: torch.Tensor, snd: torch.Tensor, case_: torch.Tensor) -> torch.Tensor:
@@ -281,12 +260,14 @@ def cond(fst: torch.Tensor, snd: torch.Tensor, case_: torch.Tensor) -> torch.Ten
     Unlike a Python conditional expression, however, the case may be a tensor,
     and both branches are evaluated, as with :func:`torch.where` ::
 
-        >>> from effectful.internals.sugar import gensym
-        >>> b = gensym(int, name="b")
-        >>> fst, snd = Indexable(torch.randn(2, 3))[b()], Indexable(torch.randn(2, 3))[b()]
+        >>> from effectful.ops.syntax import defop
+        >>> from effectful.handlers.torch import bind_dims
+
+        >>> b = defop(torch.Tensor, name="b")
+        >>> fst, snd = torch.randn(2, 3)[b()], torch.randn(2, 3)[b()]
         >>> case = (fst < snd).all(-1)
         >>> x = cond(fst, snd, case)
-        >>> assert (to_tensor(x, [b]) == to_tensor(torch.where(case[..., None], snd, fst), [b])).all()
+        >>> assert (bind_dims(x, b) == bind_dims(torch.where(case[..., None], snd, fst), b)).all()
 
     .. note::
 
@@ -304,10 +285,10 @@ def cond(fst: torch.Tensor, snd: torch.Tensor, case_: torch.Tensor) -> torch.Ten
     )
 
 
-def cond_n(values: Dict[IndexSet, torch.Tensor], case: torch.Tensor) -> torch.Tensor:
+def cond_n(values: dict[IndexSet, torch.Tensor], case: torch.Tensor) -> torch.Tensor:
     assert len(values) > 0
     assert all(isinstance(k, IndexSet) for k in values.keys())
-    result: Optional[torch.Tensor] = None
+    result: torch.Tensor | None = None
     for indices, value in values.items():
         tst = torch.as_tensor(
             functools.reduce(

effectful-0.2.0/effectful/handlers/jax/__init__.py

@@ -0,0 +1,14 @@
+try:
+    # Dummy import to check if jax is installed
+    import jax  # noqa: F401
+except ImportError:
+    raise ImportError("Jax is required to use effectful.handlers.jax")
+
+# side effect: register defdata for jax.Array
+import effectful.handlers.jax._terms  # noqa: F401
+
+from ._handlers import bind_dims as bind_dims
+from ._handlers import jax_getitem as jax_getitem
+from ._handlers import jit as jit
+from ._handlers import sizesof as sizesof
+from ._handlers import unbind_dims as unbind_dims