Change use of "opt[imize]" to "rewrite"

Author: brandonwillard

aemcmc/basic.py

@@ -5,7 +5,7 @@
 from aesara.tensor.random.utils import RandomStream
 from aesara.tensor.var import TensorVariable
 
-from aemcmc.opt import (
+from aemcmc.rewriting import (
     SamplerTracker,
     construct_ir_fgraph,
     expand_subsumptions,
@@ -41,7 +41,7 @@ def construct_sampler(
 
     fgraph.attach_feature(SamplerTracker(srng))
 
-    _ = sampler_rewrites_db.query("+basic").optimize(fgraph)
+    _ = sampler_rewrites_db.query("+basic").rewrite(fgraph)
 
     random_vars = tuple(rv for rv in fgraph.outputs if rv not in obs_rvs_to_values)
 
@@ -87,7 +87,7 @@ def construct_sampler(
         # Update the other sampled random variables in this step's graph
         sfgraph.replace_all(list(posterior_sample_steps.items()), import_missing=True)
 
-        expand_subsumptions.optimize(sfgraph)
+        expand_subsumptions.rewrite(sfgraph)
 
         step = sfgraph.outputs[0]
 

aemcmc/conjugates.py

@@ -1,13 +1,13 @@
 import aesara.tensor as at
-from aesara.graph.opt import in2out, local_optimizer
-from aesara.graph.optdb import LocalGroupDB
-from aesara.graph.unify import eval_if_etuple
+from aesara.graph.rewriting.basic import in2out, node_rewriter
+from aesara.graph.rewriting.db import LocalGroupDB
+from aesara.graph.rewriting.unify import eval_if_etuple
 from aesara.tensor.random.basic import BinomialRV
 from etuples import etuple, etuplize
 from kanren import eq, lall, run
 from unification import var
 
-from aemcmc.opt import sampler_finder_db
+from aemcmc.rewriting import sampler_finder_db
 
 
 def beta_binomial_conjugateo(observed_val, observed_rv_expr, posterior_expr):
@@ -66,7 +66,7 @@ def beta_binomial_conjugateo(observed_val, observed_rv_expr, posterior_expr):
     )
 
 
-@local_optimizer([BinomialRV])
+@node_rewriter([BinomialRV])
 def local_beta_binomial_posterior(fgraph, node):
 
     sampler_mappings = getattr(fgraph, "sampler_mappings", None)
@@ -98,7 +98,7 @@ def local_beta_binomial_posterior(fgraph, node):
     return rv_var.owner.outputs
 
 
-conjugates_db = LocalGroupDB(apply_all_opts=True)
+conjugates_db = LocalGroupDB(apply_all_rewrites=True)
 conjugates_db.name = "conjugates_db"
 conjugates_db.register("beta_binomial", local_beta_binomial_posterior, "basic")
 

aemcmc/gibbs.py

@@ -3,9 +3,9 @@
 import aesara
 import aesara.tensor as at
 from aesara.graph.basic import Variable
-from aesara.graph.opt import in2out
-from aesara.graph.optdb import LocalGroupDB
-from aesara.graph.unify import eval_if_etuple
+from aesara.graph.rewriting.basic import in2out
+from aesara.graph.rewriting.db import LocalGroupDB
+from aesara.graph.rewriting.unify import eval_if_etuple
 from aesara.ifelse import ifelse
 from aesara.tensor.math import Dot
 from aesara.tensor.random import RandomStream
@@ -19,9 +19,9 @@
     multivariate_normal_rue2005,
     polyagamma,
 )
-from aemcmc.opt import sampler_finder, sampler_finder_db
+from aemcmc.rewriting import sampler_finder, sampler_finder_db
 
-gibbs_db = LocalGroupDB(apply_all_opts=True)
+gibbs_db = LocalGroupDB(apply_all_rewrites=True)
 gibbs_db.name = "gibbs_db"
 
 

aemcmc/nuts.py

@@ -4,7 +4,11 @@
 from aehmc import nuts as aehmc_nuts
 from aehmc.utils import RaveledParamsMap
 from aeppl import joint_logprob
-from aeppl.transforms import RVTransform, TransformValuesOpt, _default_transformed_rv
+from aeppl.transforms import (
+    RVTransform,
+    TransformValuesRewrite,
+    _default_transformed_rv,
+)
 from aesara.tensor.random import RandomStream
 from aesara.tensor.var import TensorVariable
 
@@ -62,7 +66,7 @@ def nuts(
     }
 
     logprob_sum = joint_logprob(
-        model.rvs_to_values, extra_rewrites=TransformValuesOpt(transforms)
+        model.rvs_to_values, extra_rewrites=TransformValuesRewrite(transforms)
     )
 
     # Then we transform the value variables.

aemcmc/opt.py aemcmc/rewriting.pyaemcmc/opt.py renamed to aemcmc/rewriting.py

@@ -2,19 +2,19 @@
 from functools import wraps
 from typing import Callable, Dict, Iterable, List, Optional, Sequence, Set, Tuple, Union
 
-from aeppl.opt import PreserveRVMappings
+from aeppl.rewriting import PreserveRVMappings
 from aesara.compile.builders import OpFromGraph
 from aesara.compile.mode import optdb
 from aesara.graph.basic import Apply, Variable, clone_replace, io_toposort
 from aesara.graph.features import AlreadyThere, Feature
 from aesara.graph.fg import FunctionGraph
 from aesara.graph.op import Op
-from aesara.graph.opt import in2out, local_optimizer
-from aesara.graph.optdb import SequenceDB
-from aesara.tensor.basic_opt import ShapeFeature
+from aesara.graph.rewriting.basic import in2out, node_rewriter
+from aesara.graph.rewriting.db import SequenceDB
 from aesara.tensor.elemwise import DimShuffle, Elemwise
 from aesara.tensor.random.op import RandomVariable
 from aesara.tensor.random.utils import RandomStream
+from aesara.tensor.rewriting.basic import ShapeFeature
 from aesara.tensor.var import TensorVariable
 from cons.core import _car
 from unification.core import _unify
@@ -25,7 +25,7 @@
 SamplerFunctionType = Callable[
     [FunctionGraph, Apply, RandomStream], SamplerFunctionReturnType
 ]
-LocalOptimizerReturnType = Optional[Union[Dict[Variable, Variable], Sequence[Variable]]]
+LocalRewriterReturnType = Optional[Union[Dict[Variable, Variable], Sequence[Variable]]]
 
 sampler_ir_db = SequenceDB()
 sampler_ir_db.name = "sampler_ir_db"
@@ -90,7 +90,7 @@ def construct_ir_fgraph(
     # Update `obs_rvs_to_values` so that it uses the new cloned variables
     obs_rvs_to_values = {memo[k]: v for k, v in obs_rvs_to_values.items()}
 
-    sampler_ir_db.query("+basic").optimize(fgraph)
+    sampler_ir_db.query("+basic").rewrite(fgraph)
 
     new_to_old_rvs = {
         new_rv: old_rv for old_rv, new_rv in zip(rv_outputs, fgraph.outputs)
@@ -123,7 +123,7 @@ def on_attach(self, fgraph: FunctionGraph):
 
 
 def sampler_finder(tracks: Optional[Sequence[Union[Op, type]]]):
-    """Construct a `LocalOptimizer` that identifies sample steps.
+    """Construct a `NodeRewriter` that identifies sample steps.
 
     This is a decorator that is used as follows:
 
@@ -140,11 +140,11 @@ def local_horseshoe_posterior(fgraph, node, srng):
     """
 
     def decorator(f: SamplerFunctionType):
-        @local_optimizer(tracks)
+        @node_rewriter(tracks)
         @wraps(f)
         def sampler_finder(
             fgraph: FunctionGraph, node: Apply
-        ) -> LocalOptimizerReturnType:
+        ) -> LocalRewriterReturnType:
             sampler_mappings = getattr(fgraph, "sampler_mappings", None)
 
             # TODO: This assumes that `node` is a `RandomVariable`-generated `Apply` node
@@ -249,7 +249,7 @@ def car_SubsumingElemwise(x):
 _car.add((SubsumingElemwise,), car_SubsumingElemwise)
 
 
-@local_optimizer([Elemwise])
+@node_rewriter([Elemwise])
 def local_elemwise_dimshuffle_subsume(fgraph, node):
     r"""This rewrite converts `DimShuffle`s in the `Elemwise` inputs into a single `Op`.
 
@@ -359,7 +359,7 @@ def local_elemwise_dimshuffle_subsume(fgraph, node):
 )
 
 
-@local_optimizer([Elemwise])
+@node_rewriter([Elemwise])
 def inline_SubsumingElemwise(fgraph, node):
 
     op = node.op

environment.yml

@@ -11,13 +11,14 @@ dependencies:
   - compilers
   - numpy>=1.18.1
   - scipy>=1.4.0
-  - aesara>=2.6.6
-  - aeppl>=0.0.31
-  - etuples
+  - aesara>=2.8.0
+  - aeppl>=0.0.35
+  - aehmc>=0.0.9
+  - polyagamma>=1.3.2
+  - cons
   - logical-unification
+  - etuples
   - miniKanren
-  - cons
-  - polyagamma>=1.3.2
   # Intel BLAS
   - mkl
   - mkl-service

setup.py

@@ -17,9 +17,9 @@
     install_requires=[
         "numpy>=1.18.1",
         "scipy>=1.4.0",
-        "aesara>=2.6.6",
-        "aeppl>=0.0.31",
-        "aehmc>=0.0.6",
+        "aesara>=2.8.0",
+        "aeppl>=0.0.35",
+        "aehmc>=0.0.9",
         "polyagamma>=1.3.2",
         "cons",
         "logical-unification",

tests/test_basic.py

@@ -9,7 +9,7 @@
 from scipy.linalg import toeplitz
 
 from aemcmc.basic import construct_sampler
-from aemcmc.opt import SubsumingElemwise
+from aemcmc.rewriting import SubsumingElemwise
 
 
 def test_closed_form_posterior():

tests/test_conjugates.py

@@ -1,7 +1,7 @@
 import aesara
 import aesara.tensor as at
 import pytest
-from aesara.graph.unify import eval_if_etuple
+from aesara.graph.rewriting.unify import eval_if_etuple
 from aesara.tensor.random import RandomStream
 from kanren import run
 from unification import var

tests/test_gibbs.py

@@ -4,7 +4,7 @@
 import pytest
 import scipy.special
 from aesara.graph.basic import equal_computations
-from aesara.graph.opt_utils import optimize_graph
+from aesara.graph.rewriting.utils import rewrite_graph
 from aesara.tensor.random.utils import RandomStream
 from scipy.linalg import toeplitz
 
@@ -19,7 +19,7 @@
     normal_regression_posterior,
     sample_CRT,
 )
-from aemcmc.opt import SamplerTracker, construct_ir_fgraph, sampler_rewrites_db
+from aemcmc.rewriting import SamplerTracker, construct_ir_fgraph, sampler_rewrites_db
 
 
 @pytest.fixture
@@ -233,7 +233,7 @@ def test_bern_sigmoid_dot_match(srng):
     p = at.sigmoid(-eta)
     Y_rv = srng.bernoulli(p)
 
-    Y_rv = optimize_graph(Y_rv)
+    Y_rv = rewrite_graph(Y_rv)
 
     assert bern_sigmoid_dot_match(Y_rv)
 
@@ -289,7 +289,7 @@ def test_gamma_match(srng):
     b = at.scalar("b")
     beta_rv = srng.gamma(a, b)
 
-    beta_rv = optimize_graph(beta_rv)
+    beta_rv = rewrite_graph(beta_rv)
 
     a_m, b_m = gamma_match(beta_rv)
 
@@ -326,7 +326,7 @@ def test_nbinom_logistic_horseshoe_finders():
 
     fgraph.attach_feature(SamplerTracker(srng))
 
-    _ = sampler_rewrites_db.query("+basic").optimize(fgraph)
+    _ = sampler_rewrites_db.query("+basic").rewrite(fgraph)
 
     discovered_samplers = fgraph.sampler_mappings.rvs_to_samplers
     discovered_samplers = {
@@ -366,7 +366,7 @@ def test_bern_logistic_horseshoe_finders():
 
     fgraph.attach_feature(SamplerTracker(srng))
 
-    _ = sampler_rewrites_db.query("+basic").optimize(fgraph)
+    _ = sampler_rewrites_db.query("+basic").rewrite(fgraph)
 
     discovered_samplers = fgraph.sampler_mappings.rvs_to_samplers
     discovered_samplers = {