Bump PyTensor dependency and drop support for python 3.9

.github/workflows/pre-commit.yml

@@ -14,7 +14,7 @@ jobs:
     - uses: actions/checkout@9bb56186c3b09b4f86b1c65136769dd318469633
     - uses: actions/setup-python@v5
       with:
-        python-version: "3.9"  # Run pre-commit on oldest supported Python version
+        python-version: "3.10"  # Run pre-commit on oldest supported Python version
     - uses: pre-commit/action@v3.0.1
   mypy:
     runs-on: ubuntu-latest
@@ -52,7 +52,7 @@ jobs:
           activate-environment: pymc-test
           channel-priority: strict
           environment-file: conda-envs/environment-test.yml
-          python-version: "3.9"  # Run pre-commit on oldest supported Python version
+          python-version: "3.10"  # Run pre-commit on oldest supported Python version
           use-mamba: true
           use-only-tar-bz2: false # IMPORTANT: This may break caching of conda packages! See https://github.com/conda-incubator/setup-miniconda/issues/267
       - name: Install-pymc and mypy dependencies

.github/workflows/tests.yml

@@ -188,7 +188,7 @@ jobs:
       matrix:
         os: [windows-latest]
         floatx: [float64]
-        python-version: ["3.9"]
+        python-version: ["3.10"]
         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

conda-envs/environment-dev.yml

@@ -13,7 +13,7 @@ dependencies:
 - numpy>=1.15.0
 - pandas>=0.24.0
 - pip
-- pytensor>=2.19,<2.20
+- pytensor>=2.20,<2.21
 - python-graphviz
 - networkx
 - scipy>=1.4.1

conda-envs/environment-docs.yml

@@ -11,7 +11,7 @@ dependencies:
 - numpy>=1.15.0
 - pandas>=0.24.0
 - pip
-- pytensor>=2.19,<2.20
+- pytensor>=2.20,<2.21
 - python-graphviz
 - rich>=13.7.1
 - scipy>=1.4.1

conda-envs/environment-jax.yml

@@ -20,7 +20,7 @@ dependencies:
 - numpyro>=0.8.0
 - pandas>=0.24.0
 - pip
-- pytensor>=2.19,<2.20
+- pytensor>=2.20,<2.21
 - python-graphviz
 - networkx
 - rich>=13.7.1

conda-envs/environment-test.yml

@@ -16,7 +16,7 @@ dependencies:
 - numpy>=1.15.0
 - pandas>=0.24.0
 - pip
-- pytensor>=2.19,<2.20
+- pytensor>=2.20,<2.21
 - python-graphviz
 - networkx
 - rich>=13.7.1

conda-envs/windows-environment-dev.yml

@@ -13,7 +13,7 @@ dependencies:
 - numpy>=1.15.0
 - pandas>=0.24.0
 - pip
-- pytensor>=2.19,<2.20
+- pytensor>=2.20,<2.21
 - python-graphviz
 - networkx
 - rich>=13.7.1

conda-envs/windows-environment-test.yml

@@ -16,7 +16,7 @@ dependencies:
 - numpy>=1.15.0
 - pandas>=0.24.0
 - pip
-- pytensor>=2.19,<2.20
+- pytensor>=2.20,<2.21
 - python-graphviz
 - networkx
 - rich>=13.7.1

mypy.ini

@@ -1,5 +1,5 @@
 [mypy]
-python_version = 3.9
+python_version = 3.10
 no_implicit_optional = False
 strict_optional = True
 warn_redundant_casts = False

pymc/_version.py

@@ -29,7 +29,7 @@
 import subprocess
 import sys
 
-from typing import Callable
+from collections.abc import Callable
 
 
 def get_keywords():

pymc/backends/__init__.py

@@ -63,12 +63,11 @@
 
 from collections.abc import Mapping, Sequence
 from copy import copy
-from typing import Optional, Union
+from typing import Optional, TypeAlias, Union
 
 import numpy as np
 
 from pytensor.tensor.variable import TensorVariable
-from typing_extensions import TypeAlias
 
 from pymc.backends.arviz import predictions_to_inference_data, to_inference_data
 from pymc.backends.base import BaseTrace, IBaseTrace
@@ -82,7 +81,7 @@
 
     from pymc.backends.mcbackend import init_chain_adapters
 
-    TraceOrBackend = Union[BaseTrace, Backend]
+    TraceOrBackend = BaseTrace | Backend
     RunType: TypeAlias = Run
     HAS_MCB = True
 except ImportError:
@@ -98,9 +97,9 @@ def _init_trace(
     expected_length: int,
     chain_number: int,
     stats_dtypes: list[dict[str, type]],
-    trace: Optional[BaseTrace],
+    trace: BaseTrace | None,
     model: Model,
-    trace_vars: Optional[list[TensorVariable]] = None,
+    trace_vars: list[TensorVariable] | None = None,
 ) -> BaseTrace:
     """Initializes a trace backend for a chain."""
     strace: BaseTrace
@@ -119,14 +118,14 @@ def _init_trace(
 
 def init_traces(
     *,
-    backend: Optional[TraceOrBackend],
+    backend: TraceOrBackend | None,
     chains: int,
     expected_length: int,
-    step: Union[BlockedStep, CompoundStep],
+    step: BlockedStep | CompoundStep,
     initial_point: Mapping[str, np.ndarray],
     model: Model,
-    trace_vars: Optional[list[TensorVariable]] = None,
-) -> tuple[Optional[RunType], Sequence[IBaseTrace]]:
+    trace_vars: list[TensorVariable] | None = None,
+) -> tuple[RunType | None, Sequence[IBaseTrace]]:
     """Initializes a trace recorder for each chain."""
     if HAS_MCB and isinstance(backend, Backend):
         return init_chain_adapters(

pymc/backends/arviz.py

@@ -78,7 +78,7 @@ def is_data(name, var, model) -> bool:
             and var not in model.potentials
             and var not in model.value_vars
             and name not in observations
-            and isinstance(var, (Constant, SharedVariable))
+            and isinstance(var, Constant | SharedVariable)
         )
 
     # The assumption is that constants (like pm.Data) are named
@@ -163,10 +163,10 @@ def insert(self, k: str, v, idx: int):
 class InferenceDataConverter:
     """Encapsulate InferenceData specific logic."""
 
-    model: Optional[Model] = None
-    posterior_predictive: Optional[Mapping[str, np.ndarray]] = None
-    predictions: Optional[Mapping[str, np.ndarray]] = None
-    prior: Optional[Mapping[str, np.ndarray]] = None
+    model: Model | None = None
+    posterior_predictive: Mapping[str, np.ndarray] | None = None
+    predictions: Mapping[str, np.ndarray] | None = None
+    prior: Mapping[str, np.ndarray] | None = None
 
     def __init__(
         self,
@@ -177,11 +177,11 @@ def __init__(
         log_likelihood=False,
         log_prior=False,
         predictions=None,
-        coords: Optional[CoordSpec] = None,
-        dims: Optional[DimSpec] = None,
-        sample_dims: Optional[list] = None,
+        coords: CoordSpec | None = None,
+        dims: DimSpec | None = None,
+        sample_dims: list | None = None,
         model=None,
-        save_warmup: Optional[bool] = None,
+        save_warmup: bool | None = None,
         include_transformed: bool = False,
     ):
         self.save_warmup = rcParams["data.save_warmup"] if save_warmup is None else save_warmup
@@ -466,15 +466,15 @@ def to_inference_data(self):
 def to_inference_data(
     trace: Optional["MultiTrace"] = None,
     *,
-    prior: Optional[Mapping[str, Any]] = None,
-    posterior_predictive: Optional[Mapping[str, Any]] = None,
-    log_likelihood: Union[bool, Iterable[str]] = False,
-    log_prior: Union[bool, Iterable[str]] = False,
-    coords: Optional[CoordSpec] = None,
-    dims: Optional[DimSpec] = None,
-    sample_dims: Optional[list] = None,
+    prior: Mapping[str, Any] | None = None,
+    posterior_predictive: Mapping[str, Any] | None = None,
+    log_likelihood: bool | Iterable[str] = False,
+    log_prior: bool | Iterable[str] = False,
+    coords: CoordSpec | None = None,
+    dims: DimSpec | None = None,
+    sample_dims: list | None = None,
     model: Optional["Model"] = None,
-    save_warmup: Optional[bool] = None,
+    save_warmup: bool | None = None,
     include_transformed: bool = False,
 ) -> InferenceData:
     """Convert pymc data into an InferenceData object.
@@ -543,10 +543,10 @@ def predictions_to_inference_data(
     predictions,
     posterior_trace: Optional["MultiTrace"] = None,
     model: Optional["Model"] = None,
-    coords: Optional[CoordSpec] = None,
-    dims: Optional[DimSpec] = None,
-    sample_dims: Optional[list] = None,
-    idata_orig: Optional[InferenceData] = None,
+    coords: CoordSpec | None = None,
+    dims: DimSpec | None = None,
+    sample_dims: list | None = None,
+    idata_orig: InferenceData | None = None,
     inplace: bool = False,
 ) -> InferenceData:
     """Translate out-of-sample predictions into ``InferenceData``.

pymc/backends/base.py

@@ -25,9 +25,7 @@
 from collections.abc import Mapping, Sequence, Sized
 from typing import (
     Any,
-    Optional,
     TypeVar,
-    Union,
     cast,
 )
 
@@ -53,7 +51,7 @@ class IBaseTrace(ABC, Sized):
     varnames: list[str]
     """Names of tracked variables."""
 
-    sampler_vars: list[dict[str, Union[type, np.dtype]]]
+    sampler_vars: list[dict[str, type | np.dtype]]
     """Sampler stats for each sampler."""
 
     def __len__(self):
@@ -75,7 +73,7 @@ def get_values(self, varname: str, burn=0, thin=1) -> np.ndarray:
         raise NotImplementedError()
 
     def get_sampler_stats(
-        self, stat_name: str, sampler_idx: Optional[int] = None, burn=0, thin=1
+        self, stat_name: str, sampler_idx: int | None = None, burn=0, thin=1
     ) -> np.ndarray:
         """Get sampler statistics from the trace.
 
@@ -219,7 +217,7 @@ def __getitem__(self, idx):
             raise ValueError("Can only index with slice or integer")
 
     def get_sampler_stats(
-        self, stat_name: str, sampler_idx: Optional[int] = None, burn=0, thin=1
+        self, stat_name: str, sampler_idx: int | None = None, burn=0, thin=1
     ) -> np.ndarray:
         """Get sampler statistics from the trace.
 
@@ -443,7 +441,7 @@ def get_values(
         burn: int = 0,
         thin: int = 1,
         combine: bool = True,
-        chains: Optional[Union[int, Sequence[int]]] = None,
+        chains: int | Sequence[int] | None = None,
         squeeze: bool = True,
     ) -> list[np.ndarray]:
         """Get values from traces.
@@ -482,9 +480,9 @@ def get_sampler_stats(
         burn: int = 0,
         thin: int = 1,
         combine: bool = True,
-        chains: Optional[Union[int, Sequence[int]]] = None,
+        chains: int | Sequence[int] | None = None,
         squeeze: bool = True,
-    ) -> Union[list[np.ndarray], np.ndarray]:
+    ) -> list[np.ndarray] | np.ndarray:
         """Get sampler statistics from the trace.
 
         Note: This implementation attempts to squeeze object arrays into a consistent dtype,
@@ -534,7 +532,7 @@ def _slice(self, slice: slice):
         trace._report = self._report._slice(*idxs)
         return trace
 
-    def point(self, idx: int, chain: Optional[int] = None) -> dict[str, np.ndarray]:
+    def point(self, idx: int, chain: int | None = None) -> dict[str, np.ndarray]:
         """Return a dictionary of point values at `idx`.
 
         Parameters

pymc/backends/mcbackend.py

@@ -17,7 +17,7 @@
 import pickle
 
 from collections.abc import Mapping, Sequence
-from typing import Any, Optional, Union, cast
+from typing import Any, cast
 
 import hagelkorn
 import mcbackend as mcb
@@ -144,7 +144,7 @@ def _get_stats(self, fname: str, slc: slice) -> np.ndarray:
         return values
 
     def get_sampler_stats(
-        self, stat_name: str, sampler_idx: Optional[int] = None, burn=0, thin=1
+        self, stat_name: str, sampler_idx: int | None = None, burn=0, thin=1
     ) -> np.ndarray:
         slc = slice(burn, None, thin)
         # When there's just one sampler, default to remove the sampler dimension
@@ -204,7 +204,7 @@ def point(self, idx: int) -> dict[str, np.ndarray]:
 def make_runmeta_and_point_fn(
     *,
     initial_point: Mapping[str, np.ndarray],
-    step: Union[CompoundStep, BlockedStep],
+    step: CompoundStep | BlockedStep,
     model: Model,
 ) -> tuple[mcb.RunMeta, PointFunc]:
     variables, point_fn = get_variables_and_point_fn(model, initial_point)
@@ -254,7 +254,7 @@ def init_chain_adapters(
     backend: mcb.Backend,
     chains: int,
     initial_point: Mapping[str, np.ndarray],
-    step: Union[CompoundStep, BlockedStep],
+    step: CompoundStep | BlockedStep,
     model: Model,
 ) -> tuple[mcb.Run, list[ChainRecordAdapter]]:
     """Create an McBackend metadata description for the MCMC run.

pymc/backends/ndarray.py

@@ -17,7 +17,7 @@
 Store sampling values in memory as a NumPy array.
 """
 
-from typing import Any, Optional
+from typing import Any
 
 import numpy as np
 
@@ -210,7 +210,7 @@ def _slice_as_ndarray(strace, idx):
 
 
 def point_list_to_multitrace(
-    point_list: list[dict[str, np.ndarray]], model: Optional[Model] = None
+    point_list: list[dict[str, np.ndarray]], model: Model | None = None
 ) -> MultiTrace:
     """transform point list into MultiTrace"""
     _model = modelcontext(model)

pymc/backends/report.py

@@ -16,8 +16,6 @@
 import itertools
 import logging
 
-from typing import Optional
-
 from pymc.stats.convergence import _LEVELS, SamplerWarning
 
 logger = logging.getLogger(__name__)
@@ -44,17 +42,17 @@ def ok(self):
         return all(_LEVELS[warn.level] < _LEVELS["warn"] for warn in self._warnings)
 
     @property
-    def n_tune(self) -> Optional[int]:
+    def n_tune(self) -> int | None:
         """Number of tune iterations - not necessarily kept in trace!"""
         return self._n_tune
 
     @property
-    def n_draws(self) -> Optional[int]:
+    def n_draws(self) -> int | None:
         """Number of draw iterations."""
         return self._n_draws
 
     @property
-    def t_sampling(self) -> Optional[float]:
+    def t_sampling(self) -> float | None:
         """
         Number of seconds that the sampling procedure took.