Revert "Add constant_fold helper"

This reverts commit 75e9406.

Author: ricardoV94

docs/source/api/aesaraf.rst

@@ -16,7 +16,6 @@ Aesara utils
    floatX
    intX
    smartfloatX
-   constant_fold
    CallableTensor
    join_nonshared_inputs
    make_shared_replacements

pymc/aesaraf.py

@@ -34,7 +34,7 @@
 from aesara import scalar
 from aesara.compile.mode import Mode, get_mode
 from aesara.gradient import grad
-from aesara.graph import node_rewriter, rewrite_graph
+from aesara.graph import node_rewriter
 from aesara.graph.basic import (
     Apply,
     Constant,
@@ -55,13 +55,10 @@
     RandomGeneratorSharedVariable,
     RandomStateSharedVariable,
 )
-from aesara.tensor.rewriting.basic import topo_constant_folding
-from aesara.tensor.rewriting.shape import ShapeFeature
 from aesara.tensor.sharedvar import SharedVariable
 from aesara.tensor.subtensor import AdvancedIncSubtensor, AdvancedIncSubtensor1
 from aesara.tensor.var import TensorConstant, TensorVariable
 
-from pymc.exceptions import NotConstantValueError
 from pymc.vartypes import continuous_types, isgenerator, typefilter
 
 PotentialShapeType = Union[int, np.ndarray, Sequence[Union[int, Variable]], TensorVariable]
@@ -85,7 +82,6 @@
     "at_rng",
     "convert_observed_data",
     "compile_pymc",
-    "constant_fold",
 ]
 
 
@@ -975,26 +971,3 @@ def compile_pymc(
         **kwargs,
     )
     return aesara_function
-
-
-def constant_fold(xs: Sequence[TensorVariable]) -> Tuple[np.ndarray, ...]:
-    """Use constant folding to get constant values of a graph.
-
-    Parameters
-    ----------
-    xs: Sequence of TensorVariable
-        The variables that are to be constant folded
-
-    Raises
-    ------
-    NotConstantValueError:
-        If any of the variables cannot be successfully constant folded
-    """
-    fg = FunctionGraph(outputs=xs, features=[ShapeFeature()], clone=True)
-
-    folded_xs = rewrite_graph(fg, custom_rewrite=topo_constant_folding).outputs
-
-    if not all(isinstance(folded_x, Constant) for folded_x in folded_xs):
-        raise NotConstantValueError
-
-    return tuple(folded_x.data for folded_x in folded_xs)

pymc/distributions/logprob.py

@@ -27,8 +27,10 @@
 from aeppl.tensor import MeasurableJoin
 from aeppl.transforms import TransformValuesRewrite
 from aesara import tensor as at
+from aesara.graph import FunctionGraph, rewrite_graph
 from aesara.graph.basic import graph_inputs, io_toposort
 from aesara.tensor.random.op import RandomVariable
+from aesara.tensor.rewriting.basic import ShapeFeature, topo_constant_folding
 from aesara.tensor.subtensor import (
     AdvancedIncSubtensor,
     AdvancedIncSubtensor1,
@@ -39,8 +41,7 @@
 )
 from aesara.tensor.var import TensorVariable
 
-from pymc.aesaraf import constant_fold, floatX
-from pymc.exceptions import NotConstantValueError
+from pymc.aesaraf import floatX
 
 
 def _get_scaling(
@@ -337,10 +338,12 @@ def logprob_join_constant_shapes(op, values, axis, *base_vars, **kwargs):
 
     base_var_shapes = [base_var.shape[axis] for base_var in base_vars]
 
-    try:
-        base_var_shapes = constant_fold(base_var_shapes)
-    except NotConstantValueError:
-        pass
+    shape_fg = FunctionGraph(
+        outputs=base_var_shapes,
+        features=[ShapeFeature()],
+        clone=True,
+    )
+    base_var_shapes = rewrite_graph(shape_fg, custom_opt=topo_constant_folding).outputs
 
     split_values = at.split(
         value,

pymc/distributions/timeseries.py

@@ -21,12 +21,14 @@
 
 from aeppl.abstract import _get_measurable_outputs
 from aeppl.logprob import _logprob
+from aesara.graph import FunctionGraph, rewrite_graph
 from aesara.graph.basic import Node, clone_replace
 from aesara.raise_op import Assert
 from aesara.tensor import TensorVariable
 from aesara.tensor.random.op import RandomVariable
+from aesara.tensor.rewriting.basic import ShapeFeature, topo_constant_folding
 
-from pymc.aesaraf import constant_fold, convert_observed_data, floatX, intX
+from pymc.aesaraf import convert_observed_data, floatX, intX
 from pymc.distributions import distribution, multivariate
 from pymc.distributions.continuous import Flat, Normal, get_tau_sigma
 from pymc.distributions.distribution import (
@@ -44,7 +46,6 @@
     convert_dims,
     to_tuple,
 )
-from pymc.exceptions import NotConstantValueError
 from pymc.model import modelcontext
 from pymc.util import check_dist_not_registered
 
@@ -471,9 +472,14 @@ def _get_ar_order(cls, rhos: TensorVariable, ar_order: Optional[int], constant:
             If inferred ar_order cannot be inferred from rhos or if it is less than 1
         """
         if ar_order is None:
-            try:
-                (folded_shape,) = constant_fold((rhos.shape[-1],))
-            except NotConstantValueError:
+            shape_fg = FunctionGraph(
+                outputs=[rhos.shape[-1]],
+                features=[ShapeFeature()],
+                clone=True,
+            )
+            (folded_shape,) = rewrite_graph(shape_fg, custom_rewrite=topo_constant_folding).outputs
+            folded_shape = getattr(folded_shape, "data", None)
+            if folded_shape is None:
                 raise ValueError(
                     "Could not infer ar_order from last dimension of rho. Pass it "
                     "explictily or make sure rho have a static shape"

pymc/exceptions.py

@@ -78,7 +78,3 @@ def __init__(self, message, actual=None, expected=None):
 
 class TruncationError(RuntimeError):
     """Exception for errors generated from truncated graphs"""
-
-
-class NotConstantValueError(ValueError):
-    pass

pymc/tests/test_aesaraf.py

@@ -35,7 +35,6 @@
 
 from pymc.aesaraf import (
     compile_pymc,
-    constant_fold,
     convert_observed_data,
     extract_obs_data,
     replace_rng_nodes,
@@ -46,7 +45,6 @@
 from pymc.distributions.dist_math import check_parameters
 from pymc.distributions.distribution import SymbolicRandomVariable
 from pymc.distributions.transforms import Interval
-from pymc.exceptions import NotConstantValueError
 from pymc.vartypes import int_types
 
 
@@ -612,20 +610,3 @@ def test_reseed_rngs():
             assert rng.get_value()._bit_generator.state == bit_generator.state
         else:
             assert rng.get_value().bit_generator.state == bit_generator.state
-
-
-def test_constant_fold():
-    x = at.random.normal(size=(5,))
-    y = at.arange(x.size)
-
-    res = constant_fold((y, y.shape))
-    assert np.array_equal(res[0], np.arange(5))
-    assert tuple(res[1]) == (5,)
-
-
-def test_constant_fold_error():
-    x = at.vector("x")
-    y = at.arange(x.size)
-
-    with pytest.raises(NotConstantValueError):
-        constant_fold((y, y.shape))