Update staging-llm from master (#457)

* Release v0.2.3 (#374)

* Install prettyprinter for term when library is available (#386)

* install prettyprinter for term when library is available

* lint

* move code into types.py

* fix pypandoc issue (#397)

* Convert evaluate to a singledispatch (#398)

* convert evaluate to a singledispatch

* lint

* add jnp.pi and ArrayTerm.T (#394)

* Deprecate defterm (#399)

* deprecate defterm

* remove defterm case

* remove defterm

* lint

* evaluate distribution arguments

* lint

* remove interpreter

* Revert "remove interpreter"

This reverts commit 3044277.

* wip

* lint

* Rework numpyro distribution handling to enable symbolic distributions and handling of distribution methods (#311)

* refactor distribution operations

* add a test for typeof of distributions

* add tests for symbolic dists/arguments

* introduce operations for distribution methods

* comment

* fix tests

* work around #310

* replace hack with new hack

* tweak repr for _BaseOperation

* lint

* work around #312

* clean up access to dist ops

* wip

* wip

* add type annotations to get correct term conversion

* lint

* include distribution arguments as properties

* fix distribution calls

* try again

* fixes

* format

* Box the output of `__type_rule__` (#387)

* box the output of __type_rule__

* fix tests

* fix tests

* require callers of __type_rule__ to box arguments

* fix

* move Box out of ops.types

* lint

* fix test

* fix syntactic_eq implementation for jax arrays (#405)

* Fix recursion error in sizesof (#406)

* fix recursion error in sizesof

* format

* Allow `_BaseOperation` subclasses to have an overrideable `apply` method (#414)

* stash

* fixes

* initial

* wip

* lint

* ensure each subclass has a fresh operation

* wip

* wip

* lint

* wip

* wip

* lint

* refactor class method support

* move defops

* fix test

* remove singledispatch case and add test

* move definition

* cleanup

* simplify

* cleanup

* lint

* fix failing test

* fix classmethod

* __isabstractmethod__

* revert

---------

Co-authored-by: Eli <eli@basis.ai>

* Try pulling in pyproject.toml from staging-llm to master (#425)

* Generate instance-level `Operation`s for bound methods (#351)

* generalize __get__

* nits

* coverage of methoddescriptor api

* methodtype

* simplify

* simplify

* simplify

* format

* revert

* restore

* simplify

* simplify

* retain instance op on term construction

* Simplify apply inheritance

* assign

* put call next to init_subclass

* add explanatory comment

* Operation.apply -> Operation.__apply__

* add test based on issue description

* fix doctest

* Fix dataclass @defops and added dataclass metaclass (#439)

* fixed dataclass ordering and added metaclass for simplifying construction of dataclass terms

* ensure term fields are not being overriden

* added decorator and dataclass

* updated to make defdata registration automatic

* simplified dataclass loop

* updated to give property op an appropriate name

* added failing tests

* fixed failing test

* fixed numpyro/pyro/torch interfaces

* minor fix + test for deffn kwargs

* Type check and lint example code (#449)

* format example code

* type check examples

* Add beam search example using thermometer continuations (#431)

* add beam search example using thermometer continuations

* address comments

* add docstring

* lint

* Fix for jax 0.8.2 (#455)

* fix for jax 0.8.2

* add more register

* format

---------

Co-authored-by: Jack Feser <jack.feser@gmail.com>
Co-authored-by: Tim Cooijmans <cooijmans.tim@gmail.com>
Co-authored-by: Kiran Gopinathan <23038502+kiranandcode@users.noreply.github.com>

Author: 4 people

docs/source/beam.py

@@ -0,0 +1,144 @@
+"""This example demonstrates a beam search over a program that uses a `choose`
+effect for nondeterminism and `score` effect to weigh its choices.
+
+"""
+
+import functools
+import heapq
+import random
+from collections.abc import Callable
+from dataclasses import dataclass
+from pprint import pprint
+
+from effectful.ops.semantics import fwd, handler
+from effectful.ops.syntax import ObjectInterpretation, defop, implements
+
+
+@defop
+def choose[T](choices: list[T]) -> T:
+    result = random.choice(choices)
+    print(f"choose({choices}) = {result}")
+    return result
+
+
+@defop
+def score(value: float) -> None:
+    pass
+
+
+class Suspend(Exception): ...
+
+
+class ReplayIntp(ObjectInterpretation):
+    def __init__(self, trace):
+        self.trace = trace
+        self.step = 0
+
+    @implements(choose)
+    def _(self, *args, **kwargs):
+        if self.step < len(self.trace):
+            result = self.trace[self.step][1]
+            self.step += 1
+            return result
+        return fwd()
+
+
+class TraceIntp(ObjectInterpretation):
+    def __init__(self):
+        self.trace = []
+
+    @implements(choose)
+    def _(self, *args, **kwargs):
+        result = fwd()
+        self.trace.append(((args, kwargs), result))
+        return result
+
+
+class ScoreIntp(ObjectInterpretation):
+    def __init__(self):
+        self.score = 0.0
+
+    @implements(score)
+    def _(self, value):
+        self.score += value
+
+
+class ChooseOnceIntp(ObjectInterpretation):
+    def __init__(self):
+        self.is_first_call = True
+
+    @implements(choose)
+    def _(self, *args, **kwargs):
+        if not self.is_first_call:
+            raise Suspend
+
+        self.is_first_call = False
+        return fwd()
+
+
+@dataclass
+class BeamCandidate[S, T]:
+    """Represents a candidate execution path in beam search."""
+
+    trace: list[S]
+    score: float
+    in_progress: bool
+    result: T | None
+
+    def __lt__(self, other: "BeamCandidate[S, T]") -> bool:
+        return self.score < other.score
+
+    def expand[**P](self, model_fn: Callable[P, T], *args: P.args, **kwargs: P.kwargs):
+        in_progress = False
+        result = None
+        score_intp = ScoreIntp()
+        trace_intp = TraceIntp()
+        with (
+            handler(score_intp),
+            handler(ChooseOnceIntp()),
+            handler(ReplayIntp(self.trace)),
+            handler(trace_intp),
+        ):
+            try:
+                result = model_fn(*args, **kwargs)
+            except Suspend:
+                in_progress = True
+
+        return BeamCandidate(trace_intp.trace, score_intp.score, in_progress, result)
+
+
+def beam_search[**P, S, T](
+    model_fn: Callable[P, T], beam_width=3
+) -> Callable[P, BeamCandidate[S, T]]:
+    @functools.wraps(model_fn)
+    def wrapper(*args, **kwargs):
+        beam = [BeamCandidate([], 0.0, True, None)]
+
+        while True:
+            expandable = [c for c in beam if c.in_progress] * beam_width
+            if not expandable:
+                return beam
+
+            new_candidates = [c.expand(model_fn, *args, **kwargs) for c in expandable]
+
+            for c in new_candidates:
+                heapq.heappushpop(beam, c) if len(
+                    beam
+                ) >= beam_width else heapq.heappush(beam, c)
+
+    return wrapper
+
+
+if __name__ == "__main__":
+
+    def model():
+        s1 = choose(range(100))
+        score(s1)
+        s2 = choose(range(-100, 100))
+        score(s2)
+        s3 = choose(range(-100, 100))
+        score(s3)
+        return s3
+
+    result: BeamCandidate = beam_search(model)()
+    pprint(result)

docs/source/beam_search_example.rst

@@ -0,0 +1,26 @@
+Angelic Nondeterminism
+======================
+
+Here we give an example of *angelic nondeterminism* in effectful [#f1]_.
+Our model is a nondeterministic program that makes choices using a ``choose`` effect and uses a ``score`` effect to sum up a final score.
+We implement a beam search that optimizes this final score as a handler for the ``choose`` and ``score`` effects.
+
+The beam search works by running the model until it reaches a ``choose``, at which point the continuation is captured.
+This continuation is resumed multiple times with different values from ``choose`` to expand the beam.
+The intermediate score is used to rank the beam candidates.
+
+Because Python does not have support for first-class continuations, we use *thermometer continuations* [#f2]_.
+A thermometer continuation works by tracking any nondeterminism
+(essentially, the model is rerun from the start replaying the ``choose`` effects).
+If ``choose`` is the only source of nondeterminism, then the 
+after each ``choose`` and replaying it  uses *thermometer continuations* to 
+
+.. literalinclude:: ./beam.py
+    :language: python
+
+References
+----------
+
+.. [#f1] Li, Z., Solar-Lezama, A., Yue, Y., and Zheng, S., "EnCompass: Enhancing Agent Programming with Search Over Program Execution Paths", 2025. https://arxiv.org/abs/2512.03571
+
+.. [#f2] James Koppel, Gabriel Scherer, and Armando Solar-Lezama. 2018. Capturing the future by replaying the past (functional pearl). Proc. ACM Program. Lang. 2, ICFP, Article 76 (September 2018), 29 pages. https://doi.org/10.1145/3236771

docs/source/conf.py

@@ -12,10 +12,8 @@
 
 import os
 import sys
-from typing import List
 
 sys.path.insert(0, os.path.abspath("../../"))
-import sphinx_rtd_theme  # noqa: E402
 
 # -- Project information -----------------------------------------------------
 
@@ -69,7 +67,7 @@
 # List of patterns, relative to source directory, that match files and
 # directories to ignore when looking for source files.
 # This pattern also affects html_static_path and html_extra_path.
-exclude_patterns: List[str] = []
+exclude_patterns: list[str] = []
 
 
 # -- Options for HTML output -------------------------------------------------

docs/source/index.rst

@@ -16,6 +16,7 @@ Table of Contents
    minipyro_example
    lambda_example
    semi_ring_example
+   beam_search_example
 
 .. toctree::
    :maxdepth: 2

docs/source/lambda_.py

@@ -1,5 +1,6 @@
 import functools
-from typing import Annotated, Callable
+from collections.abc import Callable
+from typing import Annotated
 
 from effectful.ops.semantics import coproduct, evaluate, fvsof, fwd, handler
 from effectful.ops.syntax import Scoped, defdata, defop, syntactic_eq

docs/source/minipyro.py

@@ -14,10 +14,12 @@
 """
 
 import random
+from collections import OrderedDict
+from collections.abc import Callable
 from contextlib import contextmanager
 from dataclasses import dataclass
 from functools import partial
-from typing import Callable, Concatenate, NamedTuple, Optional, OrderedDict, Union
+from typing import Concatenate, NamedTuple
 from weakref import ref
 
 import numpy as np
@@ -40,7 +42,6 @@
 from effectful.ops.syntax import ObjectInterpretation, defop, implements
 from effectful.ops.types import Operation
 
-
 # Poutine has a notion of 'messages', which are dictionaries
 # that are passed between handlers (or 'Messengers') in order
 # to facilitate coordination and composition using "magic" slots.
@@ -56,7 +57,7 @@ class SampleMsg:
     name: str
     val: Tensor
     dist: Distribution
-    obs: Optional[Tensor]
+    obs: Tensor | None
 
 
 @dataclass
@@ -65,7 +66,7 @@ class ParamMsg:
     val: Tensor
 
 
-Message = Union[ParamMsg, SampleMsg]
+Message = ParamMsg | SampleMsg
 Trace = OrderedDict[str, Message]
 
 
@@ -87,16 +88,16 @@ class Seed(NamedTuple):
 
 
 @defop
-def sample(name: str, dist: Distribution, obs: Optional[Tensor] = None) -> Tensor:
+def sample(name: str, dist: Distribution, obs: Tensor | None = None) -> Tensor:
     raise RuntimeError("No default implementation of sample")
 
 
 @defop
 def param(
     var_name: str,
-    initial_value: Optional[Union[Tensor, Callable[[], Tensor]]] = None,
-    constraint: Optional[Constraint] = None,
-    event_dim: Optional[int] = None,
+    initial_value: Tensor | Callable[[], Tensor] | None = None,
+    constraint: Constraint | None = None,
+    event_dim: int | None = None,
 ) -> Tensor:
     raise RuntimeError("No default implementation of param")
 
@@ -123,7 +124,7 @@ def get_rng_seed() -> Seed:
 
 
 @defop
-def set_rng_seed(seed: Union[int, Seed]):
+def set_rng_seed(seed: int | Seed):
     raise RuntimeError("No default implementation of get_rng_seed")
 
 
@@ -165,9 +166,9 @@ def sample(self, var_name: str, dist: Distribution, **kwargs):
     def param(
         self,
         var_name: str,
-        initial_value: Optional[Union[Tensor, Callable[[], Tensor]]] = None,
-        constraint: Optional[Constraint] = None,
-        event_dim: Optional[int] = None,
+        initial_value: Tensor | Callable[[], Tensor] | None = None,
+        constraint: Constraint | None = None,
+        event_dim: int | None = None,
     ) -> Tensor:
         # Similar to `Tracer.sample`
 
@@ -227,7 +228,7 @@ def get_rng_seed(self):
         )
 
     @implements(set_rng_seed)
-    def set_rng_seed(self, seed: Union[int, Seed]):
+    def set_rng_seed(self, seed: int | Seed):
         if isinstance(seed, int):
             manual_seed(seed)
             random.seed(seed)
@@ -276,9 +277,9 @@ def __init__(self, initial_store=None):
     def param(
         self,
         name: str,
-        initial_value: Union[Tensor, None, Callable[[], Tensor]] = None,
+        initial_value: Tensor | None | Callable[[], Tensor] = None,
         constraint: Constraint = distributions.constraints.real,
-        event_dim: Optional[int] = None,
+        event_dim: int | None = None,
     ) -> Tensor:
         if event_dim is not None:
             raise RuntimeError("minipyro.plate does not support the event_dim arg")
@@ -321,7 +322,7 @@ class Plate(ObjectInterpretation):
     An `Interpretation` which automatically broadcasts the `sample` `Operation`
     """
 
-    def __init__(self, name: str, size: int, dim: Optional[int]):
+    def __init__(self, name: str, size: int, dim: int | None):
         if dim is None:
             raise ValueError("mini-pyro requires the `dim` argument to `plate`")
 
@@ -342,7 +343,7 @@ def do_sample(self, sampled_name: str, dist: Distribution, **kwargs) -> Tensor:
 
 
 # Helper for using `Plate` as a `handler`
-def plate(name: str, size: int, dim: Optional[int] = None):
+def plate(name: str, size: int, dim: int | None = None):
     return handler(Plate(name, size, dim))
 
 
@@ -354,7 +355,7 @@ def plate(name: str, size: int, dim: Optional[int] = None):
 
 
 def block[**P](
-    hide_fn: Callable[Concatenate[Operation, object, P], bool] = lambda *_, **__: True
+    hide_fn: Callable[Concatenate[Operation, object, P], bool] = lambda *_, **__: True,
 ):
     """
     Block is a helper for masking out a subset of calls to either
@@ -372,7 +373,7 @@ def blocking(op: Operation, *args, **kwargs):
                 return op(*args, **kwargs)
         return fwd()
 
-    return handler({sample: partial(blocking, sample), param: partial(blocking, param)})  # type: ignore
+    return handler({sample: partial(blocking, sample), param: partial(blocking, param)})
 
 
 # This is a thin wrapper around the `torch.optim.Adam` class that

docs/source/readme_example.py

@@ -35,7 +35,7 @@ def assoc_add(x, y):
 commute_rules = {add: commute_add}
 assoc_rules = {add: assoc_add}
 
-eager_mixed = functools.reduce(coproduct, (beta_rules, commute_rules, assoc_rules))
+eager_mixed = functools.reduce(coproduct, (beta_rules, commute_rules, assoc_rules))  # type: ignore
 
 x = defop(int, name="x")
 y = defop(int, name="y")

docs/source/semi_ring.py

@@ -1,7 +1,7 @@
 import collections.abc
 import operator
 import types
-from typing import Annotated, Tuple, Union, cast, overload
+from typing import Annotated, cast, overload
 
 from effectful.ops.semantics import coproduct, evaluate, fwd, handler
 from effectful.ops.syntax import Scoped, defop

effectful/handlers/jax/_terms.py

@@ -432,7 +432,8 @@ def ndim(self) -> int:
         return len(self.shape)
 
 
-@bind_dims.register  # type: ignore
+@bind_dims.register(jax.Array)  # type: ignore
+@bind_dims.register(jax._src.core.Tracer)  # type: ignore
 def _bind_dims_array(t: jax.Array, *args: Operation[[], jax.Array]) -> jax.Array:
     """Convert named dimensions to positional dimensions.
 
@@ -501,6 +502,7 @@ def _evaluate(expr):
     return reindexed
 
 
-@unbind_dims.register  # type: ignore
+@unbind_dims.register(jax.Array)  # type: ignore
+@unbind_dims.register(jax._src.core.Tracer)  # type: ignore
 def _unbind_dims_array(t: jax.Array, *args: Operation[[], jax.Array]) -> jax.Array:
     return jax_getitem(t, tuple(n() for n in args))

scripts/clean.sh

@@ -1,7 +1,7 @@
 #!/bin/bash
 set -euxo pipefail
 
-SRC="effectful tests"
+SRC="effectful tests docs/source"
 ruff check --fix $SRC
 ruff format $SRC