Update PyTensor dependency

ricardoV94 · jessegrabowski · ricardoV94 · commit b935d0d00509 · 2025-10-16T21:47:09.000+02:00
Drops support for Python 3.10 and Numpy&lt;2.0

Co-authored-by: Jesse Grabowski &lt;48652735+jessegrabowski@users.noreply.github.com&gt;
diff --git a/.github/workflows/mypy.yml b/.github/workflows/mypy.yml
@@ -19,7 +19,7 @@ jobs:
         with:
           environment-file: conda-envs/environment-test.yml
           create-args: >-
-            python=3.10
+            python=3.11
           environment-name: pymc-test
           init-shell: bash
           cache-environment: true
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
@@ -185,7 +185,7 @@ jobs:
       matrix:
         os: [windows-latest]
         floatx: [float64]
-        python-version: ["3.10"]
+        python-version: ["3.11"]
         test-subset:
           - tests/variational/test_approximations.py tests/variational/test_callbacks.py tests/variational/test_inference.py tests/variational/test_opvi.py tests/test_initial_point.py
           - tests/model/test_core.py tests/sampling/test_mcmc.py
diff --git a/conda-envs/environment-alternative-backends.yml b/conda-envs/environment-alternative-backends.yml
@@ -22,7 +22,7 @@ dependencies:
 - numpyro>=0.8.0
 - pandas>=0.24.0
 - pip
-- pytensor>=2.32.0,<2.33
+- pytensor>=2.35.0,<2.36
 - python-graphviz
 - networkx
 - rich>=13.7.1
diff --git a/conda-envs/environment-dev.yml b/conda-envs/environment-dev.yml
@@ -12,7 +12,7 @@ dependencies:
 - numpy>=1.25.0
 - pandas>=0.24.0
 - pip
-- pytensor>=2.32.0,<2.33
+- pytensor>=2.35.0,<2.36
 - python-graphviz
 - networkx
 - scipy>=1.4.1
diff --git a/conda-envs/environment-docs.yml b/conda-envs/environment-docs.yml
@@ -11,7 +11,7 @@ dependencies:
 - numpy>=1.25.0
 - pandas>=0.24.0
 - pip
-- pytensor>=2.32.0,<2.33
+- pytensor>=2.35.0,<2.36
 - python-graphviz
 - rich>=13.7.1
 - scipy>=1.4.1
diff --git a/conda-envs/environment-test.yml b/conda-envs/environment-test.yml
@@ -14,7 +14,7 @@ dependencies:
 - pandas>=0.24.0
 - pip
 - polyagamma
-- pytensor>=2.32.0,<2.33
+- pytensor>=2.35.0,<2.36
 - python-graphviz
 - networkx
 - rich>=13.7.1
diff --git a/conda-envs/windows-environment-dev.yml b/conda-envs/windows-environment-dev.yml
@@ -12,7 +12,7 @@ dependencies:
 - numpy>=1.25.0
 - pandas>=0.24.0
 - pip
-- pytensor>=2.32.0,<2.33
+- pytensor>=2.35.0,<2.36
 - python-graphviz
 - networkx
 - rich>=13.7.1
diff --git a/conda-envs/windows-environment-test.yml b/conda-envs/windows-environment-test.yml
@@ -15,7 +15,7 @@ dependencies:
 - pandas>=0.24.0
 - pip
 - polyagamma
-- pytensor>=2.32.0,<2.33
+- pytensor>=2.35.0,<2.36
 - python-graphviz
 - networkx
 - rich>=13.7.1
diff --git a/pymc/distributions/custom.py b/pymc/distributions/custom.py
@@ -18,9 +18,9 @@
 
 from pytensor import Variable, clone_replace
 from pytensor import tensor as pt
-from pytensor.graph.basic import io_toposort
 from pytensor.graph.features import ReplaceValidate
 from pytensor.graph.rewriting.basic import GraphRewriter
+from pytensor.graph.traversal import io_toposort
 from pytensor.scan.op import Scan
 from pytensor.tensor import TensorVariable, as_tensor_variable
 from pytensor.tensor.random.op import RandomVariable
diff --git a/pymc/distributions/timeseries.py b/pymc/distributions/timeseries.py
@@ -21,8 +21,9 @@
 import pytensor
 import pytensor.tensor as pt
 
-from pytensor.graph.basic import Apply, ancestors
+from pytensor.graph.basic import Apply
 from pytensor.graph.replace import clone_replace
+from pytensor.graph.traversal import ancestors
 from pytensor.tensor import TensorVariable
 from pytensor.tensor.random.op import RandomVariable
 from pytensor.tensor.random.utils import normalize_size_param
diff --git a/pymc/distributions/transforms.py b/pymc/distributions/transforms.py
@@ -17,8 +17,8 @@
 import numpy as np
 import pytensor.tensor as pt
 
+from numpy.lib.array_utils import normalize_axis_tuple
 from pytensor.graph import Op
-from pytensor.npy_2_compat import normalize_axis_tuple
 from pytensor.tensor import TensorVariable
 
 from pymc.logprob.transforms import (
diff --git a/pymc/initial_point.py b/pymc/initial_point.py
@@ -22,9 +22,10 @@
 
 from pytensor import graph_replace
 from pytensor.compile.ops import TypeCastingOp
-from pytensor.graph.basic import Apply, Variable, ancestors, walk
+from pytensor.graph.basic import Apply, Variable
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.rewriting.db import RewriteDatabaseQuery, SequenceDB
+from pytensor.graph.traversal import ancestors, walk
 from pytensor.tensor.variable import TensorVariable
 
 from pymc.logprob.transforms import Transform
diff --git a/pymc/logprob/basic.py b/pymc/logprob/basic.py
@@ -45,10 +45,9 @@
 from pytensor.graph.basic import (
     Constant,
     Variable,
-    ancestors,
-    walk,
 )
 from pytensor.graph.rewriting.basic import GraphRewriter, NodeRewriter
+from pytensor.graph.traversal import ancestors, walk
 from pytensor.tensor.variable import TensorVariable
 
 from pymc.logprob.abstract import (
@@ -533,8 +532,8 @@ def conditional_logp(
             f"The logprob terms of the following value variables could not be derived: {missing_value_terms}"
         )
 
-    values, logprobs = zip(*values_to_logprobs.items())
-    logprobs = cleanup_ir(logprobs)
+    # Ensure same order as input
+    logprobs = cleanup_ir(tuple(values_to_logprobs[v] for v in original_values))
 
     if warn_rvs:
         rvs_in_logp_expressions = _find_unallowed_rvs_in_graph(logprobs)
@@ -545,7 +544,7 @@ def conditional_logp(
                 UserWarning,
             )
 
-    return dict(zip(values, logprobs))
+    return dict(zip(original_values, logprobs))
 
 
 def transformed_conditional_logp(
diff --git a/pymc/logprob/mixture.py b/pymc/logprob/mixture.py
@@ -39,10 +39,11 @@
 import pytensor
 import pytensor.tensor as pt
 
-from pytensor.graph.basic import Apply, Constant, Variable, ancestors
+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 EquilibriumGraphRewriter, node_rewriter
+from pytensor.graph.traversal import ancestors
 from pytensor.ifelse import IfElse, ifelse
 from pytensor.scalar import Switch
 from pytensor.scalar import switch as scalar_switch
diff --git a/pymc/logprob/rewriting.py b/pymc/logprob/rewriting.py
@@ -37,11 +37,7 @@
 from collections.abc import Sequence
 
 from pytensor.compile.mode import optdb
-from pytensor.graph.basic import (
-    Variable,
-    ancestors,
-    truncated_graph_inputs,
-)
+from pytensor.graph.basic import Variable
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.replace import clone_replace
 from pytensor.graph.rewriting.basic import (
@@ -56,6 +52,7 @@
     SequenceDB,
     TopoDB,
 )
+from pytensor.graph.traversal import ancestors, truncated_graph_inputs
 from pytensor.tensor.basic import Alloc
 from pytensor.tensor.elemwise import DimShuffle, Elemwise
 from pytensor.tensor.random.rewriting import local_subtensor_rv_lift
diff --git a/pymc/logprob/tensor.py b/pymc/logprob/tensor.py
@@ -37,10 +37,10 @@
 
 from pathlib import Path
 
+from numpy.lib.array_utils import normalize_axis_index
 from pytensor import tensor as pt
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.rewriting.basic import node_rewriter
-from pytensor.npy_2_compat import normalize_axis_index
 from pytensor.tensor import TensorVariable
 from pytensor.tensor.basic import Join, MakeVector, Split
 from pytensor.tensor.elemwise import DimShuffle, Elemwise
diff --git a/pymc/logprob/utils.py b/pymc/logprob/utils.py
@@ -43,8 +43,9 @@
 
 from pytensor import tensor as pt
 from pytensor.graph import Apply, Op, node_rewriter
-from pytensor.graph.basic import Constant, clone_get_equiv, graph_inputs, walk
+from pytensor.graph.basic import Constant, clone_get_equiv
 from pytensor.graph.fg import FunctionGraph
+from pytensor.graph.traversal import graph_inputs, walk
 from pytensor.link.c.type import CType
 from pytensor.raise_op import CheckAndRaise
 from pytensor.scalar.basic import Mul
diff --git a/pymc/model/core.py b/pymc/model/core.py
@@ -34,7 +34,8 @@
 
 from pytensor.compile import DeepCopyOp, Function, ProfileStats, get_mode
 from pytensor.compile.sharedvalue import SharedVariable
-from pytensor.graph.basic import Constant, Variable, ancestors, graph_inputs
+from pytensor.graph.basic import Constant, Variable
+from pytensor.graph.traversal import ancestors, graph_inputs
 from pytensor.tensor import as_tensor
 from pytensor.tensor.math import variadic_add
 from pytensor.tensor.random.op import RandomVariable
diff --git a/pymc/model_graph.py b/pymc/model_graph.py
@@ -21,7 +21,8 @@
 from typing import Any, cast
 
 from pytensor import function
-from pytensor.graph.basic import Variable, ancestors, walk
+from pytensor.graph.basic import Variable
+from pytensor.graph.traversal import ancestors, walk
 from pytensor.tensor.shape import Shape
 
 from pymc.model.core import modelcontext
diff --git a/pymc/printing.py b/pymc/printing.py
@@ -18,7 +18,8 @@
 from functools import partial
 
 from pytensor.compile import SharedVariable
-from pytensor.graph.basic import Constant, walk
+from pytensor.graph.basic import Constant
+from pytensor.graph.traversal import walk
 from pytensor.tensor.basic import TensorVariable, Variable
 from pytensor.tensor.elemwise import DimShuffle
 from pytensor.tensor.random.basic import RandomVariable
diff --git a/pymc/pytensorf.py b/pymc/pytensorf.py
@@ -32,11 +32,10 @@
     Variable,
     clone_get_equiv,
     equal_computations,
-    graph_inputs,
-    walk,
 )
 from pytensor.graph.fg import FunctionGraph, Output
 from pytensor.graph.op import HasInnerGraph
+from pytensor.graph.traversal import graph_inputs, walk
 from pytensor.scalar.basic import Cast
 from pytensor.scan.op import Scan
 from pytensor.tensor.basic import _as_tensor_variable
diff --git a/pymc/sampling/forward.py b/pymc/sampling/forward.py
@@ -37,11 +37,9 @@
     Apply,
     Constant,
     Variable,
-    ancestors,
-    general_toposort,
-    walk,
 )
 from pytensor.graph.fg import FunctionGraph
+from pytensor.graph.traversal import ancestors, general_toposort, walk
 from pytensor.tensor.random.var import RandomGeneratorSharedVariable
 from pytensor.tensor.sharedvar import SharedVariable, TensorSharedVariable
 from pytensor.tensor.variable import TensorConstant, TensorVariable
diff --git a/pymc/sampling/jax.py b/pymc/sampling/jax.py
@@ -31,9 +31,9 @@
 from jax.lax import scan
 from numpy.typing import ArrayLike
 from pytensor.compile import SharedVariable, mode
-from pytensor.graph.basic import graph_inputs
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.replace import clone_replace
+from pytensor.graph.traversal import graph_inputs
 from pytensor.link.jax.dispatch import jax_funcify
 from pytensor.raise_op import Assert
 from pytensor.tensor import TensorVariable
diff --git a/pymc/testing.py b/pymc/testing.py
@@ -26,8 +26,9 @@
 from numpy import testing as npt
 from pytensor.compile import SharedVariable
 from pytensor.compile.mode import Mode
-from pytensor.graph.basic import Constant, Variable, equal_computations, graph_inputs
+from pytensor.graph.basic import Constant, Variable, equal_computations
 from pytensor.graph.rewriting.basic import in2out
+from pytensor.graph.traversal import graph_inputs
 from pytensor.tensor import TensorVariable
 from pytensor.tensor.random.op import RandomVariable
 from pytensor.tensor.random.type import RandomType
diff --git a/pyproject.toml b/pyproject.toml
@@ -16,7 +16,7 @@ versionfile_build = "pymc/_version.py"
 tag_prefix = "v"
 
 [tool.mypy]
-python_version = "3.10"
+python_version = "3.11"
 no_implicit_optional = false
 strict_optional = true
 warn_redundant_casts = false
@@ -30,7 +30,7 @@ warn_unused_ignores = false
 
 [tool.ruff]
 line-length = 100
-target-version = "py310"
+target-version = "py311"
 extend-exclude = ["_version.py"]
 
 [tool.ruff.format]
diff --git a/requirements-dev.txt b/requirements-dev.txt
@@ -16,7 +16,7 @@ pandas>=0.24.0
 polyagamma
 pre-commit>=2.8.0
 pymc-sphinx-theme>=0.16.0
-pytensor>=2.32.0,<2.33
+pytensor>=2.35.0,<2.36
 pytest-cov>=2.5
 pytest>=3.0
 rich>=13.7.1
diff --git a/requirements.txt b/requirements.txt
@@ -3,7 +3,7 @@ cachetools>=4.2.1
 cloudpickle
 numpy>=1.25.0
 pandas>=0.24.0
-pytensor>=2.32.0,<2.33
+pytensor>=2.35.0,<2.36
 rich>=13.7.1
 scipy>=1.4.1
 threadpoolctl>=3.1.0,<4.0.0
diff --git a/setup.py b/setup.py
@@ -30,7 +30,6 @@
     "Development Status :: 5 - Production/Stable",
     "Programming Language :: Python",
     "Programming Language :: Python :: 3",
-    "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
     "Programming Language :: Python :: 3.13",
diff --git a/tests/logprob/test_basic.py b/tests/logprob/test_basic.py
@@ -42,7 +42,8 @@
 import pytest
 import scipy.stats.distributions as sp
 
-from pytensor.graph.basic import ancestors, equal_computations
+from pytensor.graph.basic import equal_computations
+from pytensor.graph.traversal import ancestors
 from pytensor.tensor.random.op import RandomVariable
 from scipy import stats
 
@@ -89,7 +90,7 @@ def test_factorized_joint_logprob_basic():
         [ll_Y],
         rvs_to_values={sigma: sigma_value_var},
     )
-    total_ll_exp = logp(sigma, sigma_value_var) + ll_Y
+    total_ll_exp = ll_Y + logp(sigma, sigma_value_var)
 
     assert equal_computations([total_ll_combined], [total_ll_exp])
 
diff --git a/tests/logprob/test_utils.py b/tests/logprob/test_utils.py
@@ -41,7 +41,8 @@
 from pytensor import function
 from pytensor import tensor as pt
 from pytensor.compile import get_default_mode
-from pytensor.graph.basic import ancestors, equal_computations
+from pytensor.graph.basic import equal_computations
+from pytensor.graph.traversal import ancestors
 from pytensor.tensor.random.basic import NormalRV
 from pytensor.tensor.random.op import RandomVariable
 
diff --git a/tests/model/test_core.py b/tests/model/test_core.py
@@ -33,7 +33,7 @@
 import scipy.stats as st
 
 from pytensor.graph import graph_inputs
-from pytensor.graph.basic import get_var_by_name
+from pytensor.graph.traversal import get_var_by_name
 from pytensor.raise_op import Assert
 from pytensor.tensor.random.op import RandomVariable
 from pytensor.tensor.variable import TensorConstant
diff --git a/tests/sampling/test_forward.py b/tests/sampling/test_forward.py
@@ -28,7 +28,7 @@
 from pytensor import Mode, shared
 from pytensor.compile import SharedVariable
 from pytensor.graph import graph_inputs
-from pytensor.graph.basic import get_var_by_name, variable_depends_on
+from pytensor.graph.traversal import get_var_by_name, variable_depends_on
 from pytensor.tensor.variable import TensorConstant
 from scipy import stats
 
