From bcb86e06f96de0909f880c85f27bb2b953a5518c Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Thu, 4 Jul 2024 11:54:55 +0200 Subject: [PATCH 1/3] Seed flaky test --- tests/distributions/test_timeseries.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/distributions/test_timeseries.py b/tests/distributions/test_timeseries.py index 683838435a..771894737a 100644 --- a/tests/distributions/test_timeseries.py +++ b/tests/distributions/test_timeseries.py @@ -797,7 +797,7 @@ def test_batched_size(self, explicit_shape, batched_param): with Model() as t0: y = GARCH11("y", **kwargs0) - y_eval = draw(y, draws=2) + y_eval = draw(y, draws=2, random_seed=800) assert y_eval[0].shape == (batch_size, steps) assert not np.any(np.isclose(y_eval[0], y_eval[1])) From d8f3e8a9a9c1c41c83a54cfc5b3b78b79b999e0f Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Thu, 4 Jul 2024 09:58:16 +0200 Subject: [PATCH 2/3] Bump PyTensor dependency --- conda-envs/environment-dev.yml | 2 +- conda-envs/environment-docs.yml | 2 +- conda-envs/environment-jax.yml | 2 +- conda-envs/environment-test.yml | 2 +- conda-envs/windows-environment-dev.yml | 2 +- conda-envs/windows-environment-test.yml | 2 +- pymc/distributions/distribution.py | 4 ++-- pymc/logprob/order.py | 32 ++++++++++++++----------- pymc/logprob/rewriting.py | 7 ------ pymc/logprob/tensor.py | 4 +++- pymc/pytensorf.py | 4 ++-- requirements-dev.txt | 2 +- requirements.txt | 2 +- 13 files changed, 33 insertions(+), 34 deletions(-) diff --git a/conda-envs/environment-dev.yml b/conda-envs/environment-dev.yml index 7155366c33..85e6694a95 100644 --- a/conda-envs/environment-dev.yml +++ b/conda-envs/environment-dev.yml @@ -13,7 +13,7 @@ dependencies: - numpy>=1.15.0 - pandas>=0.24.0 - pip -- pytensor>=2.23,<2.24 +- pytensor>=2.25.1,<2.26 - python-graphviz - networkx - scipy>=1.4.1 diff --git a/conda-envs/environment-docs.yml b/conda-envs/environment-docs.yml index 02e28bd3c8..86097c5ab3 100644 --- a/conda-envs/environment-docs.yml +++ b/conda-envs/environment-docs.yml @@ -11,7 +11,7 @@ dependencies: - numpy>=1.15.0 - pandas>=0.24.0 - pip -- pytensor>=2.23,<2.24 +- pytensor>=2.25.1,<2.26 - python-graphviz - rich>=13.7.1 - scipy>=1.4.1 diff --git a/conda-envs/environment-jax.yml b/conda-envs/environment-jax.yml index cd2f63de23..5d79c60a8c 100644 --- a/conda-envs/environment-jax.yml +++ b/conda-envs/environment-jax.yml @@ -20,7 +20,7 @@ dependencies: - numpyro>=0.8.0 - pandas>=0.24.0 - pip -- pytensor>=2.23,<2.24 +- pytensor>=2.25.1,<2.26 - python-graphviz - networkx - rich>=13.7.1 diff --git a/conda-envs/environment-test.yml b/conda-envs/environment-test.yml index 30f685bbdb..58cde0d327 100644 --- a/conda-envs/environment-test.yml +++ b/conda-envs/environment-test.yml @@ -16,7 +16,7 @@ dependencies: - numpy>=1.15.0 - pandas>=0.24.0 - pip -- pytensor>=2.23,<2.24 +- pytensor>=2.25.1,<2.26 - python-graphviz - networkx - rich>=13.7.1 diff --git a/conda-envs/windows-environment-dev.yml b/conda-envs/windows-environment-dev.yml index 75f7ffd9e0..6d785e2cac 100644 --- a/conda-envs/windows-environment-dev.yml +++ b/conda-envs/windows-environment-dev.yml @@ -13,7 +13,7 @@ dependencies: - numpy>=1.15.0 - pandas>=0.24.0 - pip -- pytensor>=2.23,<2.24 +- pytensor>=2.25.1,<2.26 - python-graphviz - networkx - rich>=13.7.1 diff --git a/conda-envs/windows-environment-test.yml b/conda-envs/windows-environment-test.yml index b572ef1a84..fd17c31711 100644 --- a/conda-envs/windows-environment-test.yml +++ b/conda-envs/windows-environment-test.yml @@ -16,7 +16,7 @@ dependencies: - numpy>=1.15.0 - pandas>=0.24.0 - pip -- pytensor>=2.23,<2.24 +- pytensor>=2.25.1,<2.26 - python-graphviz - networkx - rich>=13.7.1 diff --git a/pymc/distributions/distribution.py b/pymc/distributions/distribution.py index d71dda97c8..51ec6ebcf1 100644 --- a/pymc/distributions/distribution.py +++ b/pymc/distributions/distribution.py @@ -836,11 +836,11 @@ def create_partial_observed_rv( if can_rewrite: masked_rv = rv[mask] fgraph = FunctionGraph(outputs=[masked_rv], clone=False, features=[ShapeFeature()]) - [unobserved_rv] = local_subtensor_rv_lift.transform(fgraph, fgraph.outputs[0].owner) + unobserved_rv = local_subtensor_rv_lift.transform(fgraph, masked_rv.owner)[masked_rv] antimasked_rv = rv[antimask] fgraph = FunctionGraph(outputs=[antimasked_rv], clone=False, features=[ShapeFeature()]) - [observed_rv] = local_subtensor_rv_lift.transform(fgraph, fgraph.outputs[0].owner) + observed_rv = local_subtensor_rv_lift.transform(fgraph, antimasked_rv.owner)[antimasked_rv] # Make a clone of the observedRV, with a distinct rng so that observed and # unobserved are never treated as equivalent (and mergeable) nodes by pytensor. diff --git a/pymc/logprob/order.py b/pymc/logprob/order.py index f15506712f..f9fa8cbe0d 100644 --- a/pymc/logprob/order.py +++ b/pymc/logprob/order.py @@ -99,18 +99,20 @@ def find_measurable_max(fgraph: FunctionGraph, node: Apply) -> list[TensorVariab if not all(params.type.broadcastable): return None - # Check whether axis covers all dimensions - axis = set(node.op.axis) - base_var_dims = set(range(base_var.ndim)) - if axis != base_var_dims: - return None + if node.op.axis is None: + axis = tuple(range(base_var.ndim)) + else: + # Check whether axis covers all dimensions + axis = tuple(sorted(node.op.axis)) + if axis != tuple(range(base_var.ndim)): + return None # distinguish measurable discrete and continuous (because logprob is different) measurable_max: Max if base_var.type.dtype.startswith("int"): - measurable_max = MeasurableMaxDiscrete(list(axis)) + measurable_max = MeasurableMaxDiscrete(axis) else: - measurable_max = MeasurableMax(list(axis)) + measurable_max = MeasurableMax(axis) max_rv_node = measurable_max.make_node(base_var) max_rv = max_rv_node.outputs @@ -206,11 +208,13 @@ def find_measurable_max_neg(fgraph: FunctionGraph, node: Apply) -> list[TensorVa if not all(params.type.broadcastable): return None - # Check whether axis is supported or not - axis = set(node.op.axis) - base_var_dims = set(range(base_var.ndim)) - if axis != base_var_dims: - return None + if node.op.axis is None: + axis = tuple(range(base_var.ndim)) + else: + # Check whether axis is supported or not + axis = tuple(sorted(node.op.axis)) + if axis != tuple(range(base_var.ndim)): + return None if not rv_map_feature.request_measurable([base_rv]): return None @@ -218,9 +222,9 @@ def find_measurable_max_neg(fgraph: FunctionGraph, node: Apply) -> list[TensorVa # distinguish measurable discrete and continuous (because logprob is different) measurable_min: Max if base_rv.type.dtype.startswith("int"): - measurable_min = MeasurableDiscreteMaxNeg(list(axis)) + measurable_min = MeasurableDiscreteMaxNeg(axis) else: - measurable_min = MeasurableMaxNeg(list(axis)) + measurable_min = MeasurableMaxNeg(axis) return measurable_min.make_node(base_rv).outputs diff --git a/pymc/logprob/rewriting.py b/pymc/logprob/rewriting.py index 76d82dd093..eb3ebd6899 100644 --- a/pymc/logprob/rewriting.py +++ b/pymc/logprob/rewriting.py @@ -72,7 +72,6 @@ from pytensor.tensor.rewriting.basic import register_canonicalize from pytensor.tensor.rewriting.math import local_exp_over_1_plus_exp from pytensor.tensor.rewriting.shape import ShapeFeature -from pytensor.tensor.rewriting.uncanonicalize import local_max_and_argmax from pytensor.tensor.subtensor import ( AdvancedIncSubtensor, AdvancedIncSubtensor1, @@ -374,12 +373,6 @@ def incsubtensor_rv_replace(fgraph, node): logprob_rewrites_db.register("measurable_ir_rewrites", measurable_ir_rewrites_db, "basic") -# Split max_and_argmax -# We only register this in the measurable IR db because max does not have a grad implemented -# And running this on any MaxAndArgmax would lead to issues: https://github.com/pymc-devs/pymc/issues/7251 -# This special registering can be removed after https://github.com/pymc-devs/pytensor/issues/334 is fixed -measurable_ir_rewrites_db.register("local_max_and_argmax", local_max_and_argmax, "basic") - # These rewrites push random/measurable variables "down", making them closer to # (or eventually) the graph outputs. Often this is done by lifting other `Op`s # "up" through the random/measurable variables and into their inputs. diff --git a/pymc/logprob/tensor.py b/pymc/logprob/tensor.py index b7c039c4b6..c709013cc6 100644 --- a/pymc/logprob/tensor.py +++ b/pymc/logprob/tensor.py @@ -35,6 +35,8 @@ # SOFTWARE. +from pathlib import Path + import pytensor from pytensor import tensor as pt @@ -237,7 +239,7 @@ class MeasurableDimShuffle(DimShuffle): # Need to get the absolute path of `c_func_file`, otherwise it tries to # find it locally and fails when a new `Op` is initialized - c_func_file = DimShuffle.get_path(DimShuffle.c_func_file) + c_func_file = str(DimShuffle.get_path(Path(DimShuffle.c_func_file))) MeasurableVariable.register(MeasurableDimShuffle) diff --git a/pymc/pytensorf.py b/pymc/pytensorf.py index f84a813963..cc7204c28a 100644 --- a/pymc/pytensorf.py +++ b/pymc/pytensorf.py @@ -36,7 +36,7 @@ graph_inputs, walk, ) -from pytensor.graph.fg import FunctionGraph +from pytensor.graph.fg import FunctionGraph, Output from pytensor.graph.op import Op from pytensor.scalar.basic import Cast from pytensor.scan.op import Scan @@ -897,7 +897,7 @@ def find_default_update(clients, rng: Variable) -> None | Variable: [client, _] = rng_clients[0] # RNG is an output of the function, this is not a problem - if client == "output": + if isinstance(client.op, Output): return rng # RNG is used by another operator, which should output an update for the RNG diff --git a/requirements-dev.txt b/requirements-dev.txt index b3e7370b1d..082eab73ce 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -17,7 +17,7 @@ numpydoc pandas>=0.24.0 polyagamma pre-commit>=2.8.0 -pytensor>=2.23,<2.24 +pytensor>=2.25.1,<2.26 pytest-cov>=2.5 pytest>=3.0 rich>=13.7.1 diff --git a/requirements.txt b/requirements.txt index c330cd56dd..b59ca29127 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,7 +3,7 @@ cachetools>=4.2.1 cloudpickle numpy>=1.15.0 pandas>=0.24.0 -pytensor>=2.23,<2.24 +pytensor>=2.25.1,<2.26 rich>=13.7.1 scipy>=1.4.1 threadpoolctl>=3.1.0,<4.0.0 From 15f139edcaf5804d261421faa7e80ae8b64bfa4b Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Tue, 9 Jul 2024 16:09:00 +0200 Subject: [PATCH 3/3] Test less extreme combinations of parameters in SkewStudentT logcdf Scipy underflows earlier than PyTensor in newer versions --- tests/distributions/test_continuous.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/distributions/test_continuous.py b/tests/distributions/test_continuous.py index e68de7732b..73ab07d3c1 100644 --- a/tests/distributions/test_continuous.py +++ b/tests/distributions/test_continuous.py @@ -592,7 +592,7 @@ def test_skewstudentt_logcdf(self): check_logcdf( pm.SkewStudentT, R, - {"a": Rplus, "b": Rplus, "mu": R, "sigma": Rplus}, + {"a": Rplus, "b": Rplus, "mu": R, "sigma": Rplusbig}, lambda value, a, b, mu, sigma: st.jf_skew_t.logcdf(value, a, b, mu, sigma), )