diff --git a/conda-envs/environment-dev.yml b/conda-envs/environment-dev.yml index 66595e81820..5b6b98c936e 100644 --- a/conda-envs/environment-dev.yml +++ b/conda-envs/environment-dev.yml @@ -14,7 +14,7 @@ dependencies: - numpy>=1.15.0 - pandas>=0.24.0 - pip -- pytensor>=2.12.0,<2.13 +- pytensor>=2.13.0,<2.14 - python-graphviz - networkx - scipy>=1.4.1 diff --git a/conda-envs/environment-docs.yml b/conda-envs/environment-docs.yml index b1c2b4ee7f5..87a6bf315bb 100644 --- a/conda-envs/environment-docs.yml +++ b/conda-envs/environment-docs.yml @@ -12,7 +12,7 @@ dependencies: - numpy>=1.15.0 - pandas>=0.24.0 - pip -- pytensor>=2.12.0,<2.13 +- pytensor>=2.13.0,<2.14 - python-graphviz - scipy>=1.4.1 - typing-extensions>=3.7.4 diff --git a/conda-envs/environment-test.yml b/conda-envs/environment-test.yml index 4ce5e06a6df..0dd7f8efc79 100644 --- a/conda-envs/environment-test.yml +++ b/conda-envs/environment-test.yml @@ -17,7 +17,7 @@ dependencies: - numpy>=1.15.0 - pandas>=0.24.0 - pip -- pytensor>=2.12.0,<2.13 +- pytensor>=2.13.0,<2.14 - python-graphviz - networkx - scipy>=1.4.1 diff --git a/conda-envs/windows-environment-dev.yml b/conda-envs/windows-environment-dev.yml index 86d1830d709..bf64ee8f4dd 100644 --- a/conda-envs/windows-environment-dev.yml +++ b/conda-envs/windows-environment-dev.yml @@ -14,7 +14,7 @@ dependencies: - numpy>=1.15.0 - pandas>=0.24.0 - pip -- pytensor>=2.12.0,<2.13 +- pytensor>=2.13.0,<2.14 - python-graphviz - networkx - scipy>=1.4.1 diff --git a/conda-envs/windows-environment-test.yml b/conda-envs/windows-environment-test.yml index 15953058364..a41f130b60d 100644 --- a/conda-envs/windows-environment-test.yml +++ b/conda-envs/windows-environment-test.yml @@ -17,7 +17,7 @@ dependencies: - numpy>=1.15.0 - pandas>=0.24.0 - pip -- pytensor>=2.12.0,<2.13 +- pytensor>=2.13.0,<2.14 - python-graphviz - networkx - scipy>=1.4.1 diff --git a/pymc/distributions/distribution.py b/pymc/distributions/distribution.py index 5c74080ace1..252268a728f 100644 --- a/pymc/distributions/distribution.py +++ b/pymc/distributions/distribution.py @@ -34,6 +34,7 @@ from pytensor.tensor.random.op import RandomVariable from pytensor.tensor.random.rewriting import local_subtensor_rv_lift from pytensor.tensor.random.utils import normalize_size_param +from pytensor.tensor.rewriting.shape import ShapeFeature from pytensor.tensor.var import TensorVariable from typing_extensions import TypeAlias @@ -54,7 +55,12 @@ from pymc.logprob.rewriting import logprob_rewrites_db from pymc.model import new_or_existing_block_model_access from pymc.printing import str_for_dist -from pymc.pytensorf import collect_default_updates, convert_observed_data, floatX +from pymc.pytensorf import ( + collect_default_updates, + constant_fold, + convert_observed_data, + floatX, +) from pymc.util import UNSET, _add_future_warning_tag from pymc.vartypes import continuous_types, string_types @@ -1229,15 +1235,12 @@ def create_partial_observed_rv( can_rewrite = True if can_rewrite: - # Rewrite doesn't work with boolean masks. Should be fixed after https://github.com/pymc-devs/pytensor/pull/329 - mask, antimask = mask.nonzero(), antimask.nonzero() - masked_rv = rv[mask] - fgraph = FunctionGraph(outputs=[masked_rv], clone=False) + fgraph = FunctionGraph(outputs=[masked_rv], clone=False, features=[ShapeFeature()]) [unobserved_rv] = local_subtensor_rv_lift.transform(fgraph, fgraph.outputs[0].owner) antimasked_rv = rv[antimask] - fgraph = FunctionGraph(outputs=[antimasked_rv], clone=False) + fgraph = FunctionGraph(outputs=[antimasked_rv], clone=False, features=[ShapeFeature()]) [observed_rv] = local_subtensor_rv_lift.transform(fgraph, fgraph.outputs[0].owner) # Make a clone of the observedRV, with a distinct rng so that observed and @@ -1270,7 +1273,7 @@ def partial_observed_rv_logprob(op, values, dist, mask, **kwargs): # For the logp, simply join the values [obs_value, unobs_value] = values antimask = ~mask - joined_value = pt.empty_like(dist) + joined_value = pt.empty(constant_fold([dist.shape])[0]) joined_value = pt.set_subtensor(joined_value[mask], unobs_value) joined_value = pt.set_subtensor(joined_value[antimask], obs_value) joined_logp = logp(dist, joined_value) diff --git a/pymc/gp/cov.py b/pymc/gp/cov.py index 9fc5e62affd..c5ffcac1d9f 100644 --- a/pymc/gp/cov.py +++ b/pymc/gp/cov.py @@ -183,7 +183,10 @@ def n_dims(self) -> int: def _slice(self, X, Xs=None): xdims = X.shape[-1] if isinstance(xdims, Variable): - xdims = xdims.eval() + # Circular dependency + from pymc.pytensorf import constant_fold + + xdims = constant_fold([xdims])[0] if self.input_dim != xdims: warnings.warn( f"Only {self.input_dim} column(s) out of {xdims} are" diff --git a/pymc/logprob/mixture.py b/pymc/logprob/mixture.py index 267670a835e..daa16f520e6 100644 --- a/pymc/logprob/mixture.py +++ b/pymc/logprob/mixture.py @@ -42,7 +42,7 @@ from pytensor.graph.basic import Apply, Constant, Variable from pytensor.graph.fg import FunctionGraph from pytensor.graph.op import Op, compute_test_value -from pytensor.graph.rewriting.basic import node_rewriter, pre_greedy_node_rewriter +from pytensor.graph.rewriting.basic import EquilibriumGraphRewriter, node_rewriter from pytensor.ifelse import IfElse, ifelse from pytensor.scalar import Switch from pytensor.scalar import switch as scalar_switch @@ -52,6 +52,7 @@ local_rv_size_lift, local_subtensor_rv_lift, ) +from pytensor.tensor.rewriting.shape import ShapeFeature from pytensor.tensor.shape import shape_tuple from pytensor.tensor.subtensor import ( AdvancedSubtensor, @@ -77,7 +78,6 @@ measurable_ir_rewrites_db, subtensor_ops, ) -from pymc.logprob.tensor import naive_bcast_rv_lift from pymc.logprob.utils import check_potential_measurability @@ -203,21 +203,17 @@ def expand_indices( return cast(Tuple[TensorVariable], tuple(pt.broadcast_arrays(*adv_indices))) -def rv_pull_down(x: TensorVariable, dont_touch_vars=None) -> TensorVariable: +def rv_pull_down(x: TensorVariable) -> TensorVariable: """Pull a ``RandomVariable`` ``Op`` down through a graph, when possible.""" - fgraph = FunctionGraph(outputs=dont_touch_vars or [], clone=False) - - return pre_greedy_node_rewriter( - fgraph, - [ - local_rv_size_lift, - local_dimshuffle_rv_lift, - local_subtensor_rv_lift, - naive_bcast_rv_lift, - local_lift_DiracDelta, - ], - x, - ) + fgraph = FunctionGraph(outputs=[x], clone=False, features=[ShapeFeature()]) + rewrites = [ + local_rv_size_lift, + local_dimshuffle_rv_lift, + local_subtensor_rv_lift, + local_lift_DiracDelta, + ] + EquilibriumGraphRewriter(rewrites, max_use_ratio=100).rewrite(fgraph) + return fgraph.outputs[0] class MixtureRV(Op): diff --git a/pymc/math.py b/pymc/math.py index 2f8520b0ecc..0cb86095ee4 100644 --- a/pymc/math.py +++ b/pymc/math.py @@ -75,13 +75,14 @@ where, zeros_like, ) +from pytensor.tensor.special import log_softmax, softmax try: from pytensor.tensor.basic import extract_diag except ImportError: from pytensor.tensor.nlinalg import extract_diag -from pytensor.tensor.nlinalg import det, matrix_dot, matrix_inverse, trace +from pytensor.tensor.nlinalg import matrix_inverse from scipy.linalg import block_diag as scipy_block_diag from pymc.pytensorf import floatX, ix_, largest_common_dtype @@ -267,31 +268,7 @@ def logdiffexp_numpy(a, b): return a + log1mexp_numpy(b - a, negative_input=True) -def invlogit(x, eps=None): - """The inverse of the logit function, 1 / (1 + exp(-x)).""" - if eps is not None: - warnings.warn( - "pymc.math.invlogit no longer supports the ``eps`` argument and it will be ignored.", - FutureWarning, - stacklevel=2, - ) - return pt.sigmoid(x) - - -def softmax(x, axis=None): - # Ignore vector case UserWarning issued by PyTensor. This can be removed once PyTensor - # drops that warning - with warnings.catch_warnings(): - warnings.simplefilter("ignore", UserWarning) - return pt.special.softmax(x, axis=axis) - - -def log_softmax(x, axis=None): - # Ignore vector case UserWarning issued by PyTensor. This can be removed once PyTensor - # drops that warning - with warnings.catch_warnings(): - warnings.simplefilter("ignore", UserWarning) - return pt.special.log_softmax(x, axis=axis) +invlogit = sigmoid def logbern(log_p): diff --git a/pymc/model_graph.py b/pymc/model_graph.py index 43e7d1f728d..1f3e6d1b895 100644 --- a/pymc/model_graph.py +++ b/pymc/model_graph.py @@ -23,6 +23,7 @@ from pytensor.scalar.basic import Cast from pytensor.tensor.elemwise import Elemwise from pytensor.tensor.random.op import RandomVariable +from pytensor.tensor.shape import Shape from pytensor.tensor.var import TensorConstant, TensorVariable import pymc as pm @@ -55,6 +56,9 @@ def get_parent_names(self, var: TensorVariable) -> Set[VarName]: def _filter_non_parameter_inputs(var): node = var.owner + if isinstance(node.op, Shape): + # Don't show shape-related dependencies + return [] if isinstance(node.op, RandomVariable): # Filter out rng, dtype and size parameters or RandomVariable nodes return node.inputs[3:] diff --git a/pymc/pytensorf.py b/pymc/pytensorf.py index 04ed1bdad24..066ba233c5b 100644 --- a/pymc/pytensorf.py +++ b/pymc/pytensorf.py @@ -87,7 +87,6 @@ "generator", "convert_observed_data", "compile_pymc", - "constant_fold", ] diff --git a/requirements-dev.txt b/requirements-dev.txt index 6032958e14e..713c09ceeb2 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -18,7 +18,7 @@ numpydoc pandas>=0.24.0 polyagamma pre-commit>=2.8.0 -pytensor>=2.12.0,<2.13 +pytensor>=2.13.0,<2.14 pytest-cov>=2.5 pytest>=3.0 scipy>=1.4.1 diff --git a/requirements.txt b/requirements.txt index 82d327e62a3..41c350234e4 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,6 +4,6 @@ cloudpickle fastprogress>=0.2.0 numpy>=1.15.0 pandas>=0.24.0 -pytensor>=2.12.0,<2.13 +pytensor>=2.13.0,<2.14 scipy>=1.4.1 typing-extensions>=3.7.4 diff --git a/tests/distributions/test_dist_math.py b/tests/distributions/test_dist_math.py index f0b6d8f87f7..2ccab76b1f1 100644 --- a/tests/distributions/test_dist_math.py +++ b/tests/distributions/test_dist_math.py @@ -224,7 +224,7 @@ def check_vals(fn1, fn2, *args): def test_multigamma(): - x = pt.vector("x") + x = pt.vector("x", shape=(1,)) p = pt.scalar("p") xvals = [np.array([v], dtype=config.floatX) for v in [0.1, 2, 5, 10, 50, 100]] diff --git a/tests/logprob/test_mixture.py b/tests/logprob/test_mixture.py index 53e2096b554..ef46b2cf27a 100644 --- a/tests/logprob/test_mixture.py +++ b/tests/logprob/test_mixture.py @@ -1119,6 +1119,7 @@ def test_ifelse_mixture_shared_component(): ) +@pytest.mark.xfail(reason="Relied on rewrite-case that is no longer supported by PyTensor") def test_joint_logprob_subtensor(): """Make sure we can compute a joint log-probability for ``Y[I]`` where ``Y`` and ``I`` are random variables.""" @@ -1137,6 +1138,10 @@ def test_joint_logprob_subtensor(): I_rv = pt.random.bernoulli(p, size=size, rng=rng) I_rv.name = "I" + # The rewrite for lifting subtensored RVs refuses to work with advanced + # indexing as it could lead to repeated draws. + # TODO: Re-enable rewrite for cases where this is not a concern + # (e.g., at least one of the advanced indexes has non-repeating values) A_idx = A_rv[I_rv, pt.ogrid[A_rv.shape[-1] :]] assert isinstance(A_idx.owner.op, (Subtensor, AdvancedSubtensor, AdvancedSubtensor1)) diff --git a/tests/test_math.py b/tests/test_math.py index c5aed9246cf..ff9c8708c00 100644 --- a/tests/test_math.py +++ b/tests/test_math.py @@ -262,37 +262,3 @@ def test_expand_packed_triangular(): assert np.all(expand_upper.eval({packed: upper_packed}) == upper) assert np.all(expand_diag_lower.eval({packed: lower_packed}) == floatX(np.diag(vals))) assert np.all(expand_diag_upper.eval({packed: upper_packed}) == floatX(np.diag(vals))) - - -def test_invlogit_deprecation_warning(): - with pytest.warns( - FutureWarning, - match="pymc.math.invlogit no longer supports the", - ): - res = invlogit(np.array(-750.0), 1e-5).eval() - - with warnings.catch_warnings(): - warnings.simplefilter("error") - res_zero_eps = invlogit(np.array(-750.0)).eval() - - assert np.isclose(res, res_zero_eps) - - -@pytest.mark.parametrize( - "pytensor_function, pymc_wrapper", - [ - (pt.special.softmax, softmax), - (pt.special.log_softmax, log_softmax), - ], -) -def test_softmax_logsoftmax_no_warnings(pytensor_function, pymc_wrapper): - """Test that wrappers for pytensor functions do not issue Warnings""" - - vector = pt.vector("vector") - with pytest.warns(Warning) as record: - pytensor_function(vector) - assert {w.category for w in record.list} == {UserWarning, FutureWarning} - - with warnings.catch_warnings(): - warnings.simplefilter("error") - pymc_wrapper(vector) diff --git a/tests/test_model.py b/tests/test_model.py index f4d2bbe78d1..c4d04918a07 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -963,7 +963,7 @@ def test_set_data_constant_shape_error(): pmodel.add_coord("weekday", length=x.shape[0]) pm.MutableData("y", np.arange(7), dims="weekday") - msg = "because the dimension length is tied to a TensorConstant" + msg = "because the dimension was initialized from 'x'" with pytest.raises(ShapeError, match=msg): pmodel.set_data("y", np.arange(10))