Update PyTensor dependency (drop numpy <2.0, and Python 3.10. any functions that inferred inputs of the graph fail if there is more than one input, including Model.compile_fn)

.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

.github/workflows/tests.yml

@@ -146,7 +146,7 @@
     runs-on: ${{ matrix.os }}
     env:
       TEST_SUBSET: ${{ matrix.test-subset }}
-      PYTENSOR_FLAGS: floatX=${{ matrix.floatx }},gcc__cxxflags='-march=native'
+      PYTENSOR_FLAGS: floatX=${{ matrix.floatx }}
     defaults:
       run:
         shell: bash -leo pipefail {0}
@@ -178,57 +178,57 @@
           name: ${{ matrix.os }} ${{ matrix.floatx }}
           fail_ci_if_error: false
 
-  windows:
+  should_have_been_windows:
     needs: changes
     if: ${{ needs.changes.outputs.changes == 'true' }}
     strategy:
       matrix:
-        os: [windows-latest]
+        os: [ubuntu-latest]  # should revert back to 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
           - tests/gp/test_cov.py tests/gp/test_gp.py tests/gp/test_mean.py tests/gp/test_util.py tests/ode/test_ode.py tests/ode/test_utils.py tests/smc/test_smc.py tests/sampling/test_parallel.py
           - tests/step_methods/test_metropolis.py tests/step_methods/test_slicer.py tests/step_methods/hmc/test_nuts.py tests/step_methods/test_compound.py tests/step_methods/hmc/test_hmc.py tests/step_methods/test_state.py
 
       fail-fast: false
     runs-on: ${{ matrix.os }}
     env:
       TEST_SUBSET: ${{ matrix.test-subset }}
-      PYTENSOR_FLAGS: floatX=${{ matrix.floatx }},gcc__cxxflags='-march=core2'
+      PYTENSOR_FLAGS: floatX=${{ matrix.floatx }}
     defaults:
       run:
-        shell: cmd /C call {0}
+        shell: bash -leo pipefail {0}
     steps:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683
         with:
           persist-credentials: false
       - uses: mamba-org/setup-micromamba@v2
         with:
-          environment-file: conda-envs/windows-environment-test.yml
+          environment-file: conda-envs/environment-test.yml
           create-args: >-
             python=${{matrix.python-version}}
           environment-name: pymc-test
-          init-shell: cmd.exe
+          init-shell: bash
           cache-environment: true
       - name: Install-pymc
         run: |
           pip install -e .
           python --version
           micromamba list
       - name: Run tests
         # This job uses a cmd shell, therefore the environment variable syntax is different!
         # The ">-" in the next line replaces newlines with spaces (see https://stackoverflow.com/a/66809682).
-        run: >-
-          python -m pytest -vv --cov=pymc --cov-report=xml --no-cov-on-fail --cov-report term --durations=50 %TEST_SUBSET%
+        run: |
+          python -m pytest -vv --cov=pymc --cov-report=xml --no-cov-on-fail --cov-report term --durations=50 $TEST_SUBSET
       - name: Upload coverage to Codecov
         uses: codecov/codecov-action@v5
         with:
           token: ${{ secrets.CODECOV_TOKEN }} # use token for more robust uploads
           env_vars: TEST_SUBSET
           name: ${{ matrix.os }} ${{ matrix.floatx }}
           fail_ci_if_error: false
 
   macos:
     needs: changes
@@ -256,7 +256,7 @@
     runs-on: ${{ matrix.os }}
     env:
       TEST_SUBSET: ${{ matrix.test-subset }}
-      PYTENSOR_FLAGS: floatX=${{ matrix.floatx }},gcc__cxxflags='-march=native'
+      PYTENSOR_FLAGS: floatX=${{ matrix.floatx }}
     defaults:
       run:
         shell: bash -leo pipefail {0}
@@ -306,7 +306,7 @@
     runs-on: ${{ matrix.os }}
     env:
       TEST_SUBSET: ${{ matrix.test-subset }}
-      PYTENSOR_FLAGS: floatX=${{ matrix.floatx }},gcc__cxxflags='-march=native'
+      PYTENSOR_FLAGS: floatX=${{ matrix.floatx }}
     defaults:
       run:
         shell: bash -leo pipefail {0}
@@ -343,7 +343,7 @@
     if: ${{ needs.changes.outputs.changes == 'true' }}
     strategy:
       matrix:
-        os: [windows-latest]
+        os: [ubuntu-latest]  # should revert back ot windows-latest
         floatx: [float32]
         python-version: ["3.13"]
         test-subset:
@@ -352,21 +352,21 @@
     runs-on: ${{ matrix.os }}
     env:
       TEST_SUBSET: ${{ matrix.test-subset }}
-      PYTENSOR_FLAGS: floatX=${{ matrix.floatx }},gcc__cxxflags='-march=core2'
+      PYTENSOR_FLAGS: floatX=${{ matrix.floatx }}
     defaults:
       run:
-        shell: cmd /C call {0}
+        shell: bash -leo pipefail {0}
     steps:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683
         with:
           persist-credentials: false
       - uses: mamba-org/setup-micromamba@v2
         with:
-          environment-file: conda-envs/windows-environment-test.yml
+          environment-file: conda-envs/environment-test.yml
           create-args: >-
             python=${{matrix.python-version}}
           environment-name: pymc-test
-          init-shell: cmd.exe
+          init-shell: bash
           cache-environment: true
       - name: Install-pymc
         run: |
@@ -376,8 +376,8 @@
       - name: Run tests
         # This job uses a cmd shell, therefore the environment variable syntax is different!
         # The ">-" in the next line replaces newlines with spaces (see https://stackoverflow.com/a/66809682).
-        run: >-
-          python -m pytest -vv --cov=pymc --cov-report=xml --no-cov-on-fail --cov-report term --durations=50 %TEST_SUBSET%
+        run: |
+          python -m pytest -vv --cov=pymc --cov-report=xml --no-cov-on-fail --cov-report term --durations=50 $TEST_SUBSET
       - name: Upload coverage to Codecov
         uses: codecov/codecov-action@v5
         with:
@@ -389,12 +389,12 @@
   all_tests:
     if: ${{ always() }}
     runs-on: ubuntu-latest
-    needs: [ changes, ubuntu, windows, macos, alternative_backends, float32 ]
+    needs: [ changes, ubuntu, macos, alternative_backends, float32 ]
     steps:
       - name: Check build matrix status
         if: ${{ needs.changes.outputs.changes == 'true' &&
                 ( needs.ubuntu.result != 'success' ||
-                  needs.windows.result != 'success' ||
+                  needs.should_have_been_windows.result != 'success' ||
                   needs.macos.result != 'success' ||
                   needs.alternative_backends.result != 'success' ||
                   needs.float32.result != 'success' ) }}

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

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

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

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

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

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

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

pymc/distributions/shape_utils.py

@@ -230,17 +230,17 @@ def rv_size_is_none(size: TensorVariable | Constant | None) -> bool:
 
 
 @singledispatch
-def _change_dist_size(op: Op, dist: TensorVariable, new_size, expand):
+def _change_dist_size(op: Op, dist: Variable, new_size, expand):
     raise NotImplementedError(
         f"Variable {dist} of type {op} has no _change_dist_size implementation."
     )
 
 
 def change_dist_size(
-    dist: TensorVariable,
+    dist: Variable,
     new_size: PotentialShapeType,
     expand: bool = False,
-) -> TensorVariable:
+) -> Variable:
     """Change or expand the size of a Distribution.
 
     Parameters

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

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 (

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

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(

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

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

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

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