Add constant_fold helper

Author: ricardoV94

docs/source/api/aesaraf.rst

@@ -16,6 +16,7 @@ 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
+from aesara.graph import node_rewriter, rewrite_graph
 from aesara.graph.basic import (
     Apply,
     Constant,
@@ -55,10 +55,13 @@
     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]
@@ -82,6 +85,7 @@
     "at_rng",
     "convert_observed_data",
     "compile_pymc",
+    "constant_fold",
 ]
 
 
@@ -971,3 +975,30 @@ def compile_pymc(
         **kwargs,
     )
     return aesara_function
+
+
+def constant_fold(
+    xs: Sequence[TensorVariable], raise_not_constant: bool = True
+) -> 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
+    raise_not_constant: bool, default True
+        Raises NotConstantValueError if any of the variables cannot be constant folded.
+        This should only be disabled with care, as the graphs are cloned before
+        attempting constant folding, and any old non-shared inputs will not work with
+        the returned outputs
+    """
+    fg = FunctionGraph(outputs=xs, features=[ShapeFeature()], clone=True)
+
+    folded_xs = rewrite_graph(fg, custom_rewrite=topo_constant_folding).outputs
+
+    if raise_not_constant and not all(isinstance(folded_x, Constant) for folded_x in folded_xs):
+        raise NotConstantValueError
+
+    return tuple(
+        folded_x.data if isinstance(folded_x, Constant) else folded_x for folded_x in folded_xs
+    )

pymc/distributions/logprob.py

@@ -27,10 +27,8 @@
 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,
@@ -41,7 +39,7 @@
 )
 from aesara.tensor.var import TensorVariable
 
-from pymc.aesaraf import floatX
+from pymc.aesaraf import constant_fold, floatX
 
 
 def _get_scaling(
@@ -338,12 +336,8 @@ def logprob_join_constant_shapes(op, values, axis, *base_vars, **kwargs):
 
     base_var_shapes = [base_var.shape[axis] for base_var in base_vars]
 
-    shape_fg = FunctionGraph(
-        outputs=base_var_shapes,
-        features=[ShapeFeature()],
-        clone=True,
-    )
-    base_var_shapes = rewrite_graph(shape_fg, custom_opt=topo_constant_folding).outputs
+    # We don't need the graph to be constant, just to have RandomVariables removed
+    base_var_shapes = constant_fold(base_var_shapes, raise_not_constant=False)
 
     split_values = at.split(
         value,

pymc/distributions/timeseries.py

@@ -21,14 +21,12 @@
 
 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 convert_observed_data, floatX, intX
+from pymc.aesaraf import constant_fold, 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 (
@@ -46,6 +44,7 @@
     convert_dims,
     to_tuple,
 )
+from pymc.exceptions import NotConstantValueError
 from pymc.model import modelcontext
 from pymc.util import check_dist_not_registered
 
@@ -472,14 +471,9 @@ 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:
-            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:
+            try:
+                (folded_shape,) = constant_fold((rhos.shape[-1],))
+            except NotConstantValueError:
                 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,3 +78,7 @@ 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,6 +35,7 @@
 
 from pymc.aesaraf import (
     compile_pymc,
+    constant_fold,
     convert_observed_data,
     extract_obs_data,
     replace_rng_nodes,
@@ -45,6 +46,7 @@
 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
 
 
@@ -610,3 +612,24 @@ 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_raises():
+    size = aesara.shared(5)
+    x = at.random.normal(size=(size,))
+    y = at.arange(x.size)
+
+    with pytest.raises(NotConstantValueError):
+        constant_fold((y, y.shape))
+
+    res = constant_fold((y, y.shape), raise_not_constant=False)
+    assert tuple(res[1].eval()) == (5,)