Bump PyTensor and support numpy>2.0 and Python 3.13

.github/workflows/mypy.yml

@@ -26,7 +26,6 @@ jobs:
       - name: Install-pymc and mypy dependencies
         run: |
           pip install -e .
-          pip install --pre -U polyagamma
           python --version
       - name: Run mypy
         run: |

.github/workflows/tests.yml

@@ -56,7 +56,7 @@ jobs:
       matrix:
         os: [ubuntu-20.04]
         floatx: [float64]
-        python-version: ["3.12"]
+        python-version: ["3.13"]
         test-subset:
           - |
             tests/test_util.py
@@ -158,8 +158,6 @@ jobs:
       - name: Install-pymc
         run: |
           pip install -e .
-          # TODO: https://github.com/pymc-devs/pymc/issues/7417
-          pip install --pre -U 'polyagamma<1.3.7'
           python --version
           micromamba list
       - name: Run tests
@@ -210,7 +208,6 @@ jobs:
       - name: Install-pymc
         run: |
           pip install -e .
-          pip install --pre -U 'polyagamma<1.3.7'
           python --version
           micromamba list
       - name: Run tests
@@ -233,7 +230,7 @@ jobs:
       matrix:
         os: [macos-latest]
         floatx: [float64]
-        python-version: ["3.12"]
+        python-version: ["3.13"]
         test-subset:
           - |
             tests/sampling/test_parallel.py
@@ -291,7 +288,7 @@ jobs:
       matrix:
         os: [ubuntu-20.04]
         floatx: [float64]
-        python-version: ["3.12"]
+        python-version: ["3.13"]
         test-subset:
           - tests/sampling/test_jax.py tests/sampling/test_mcmc_external.py
       fail-fast: false
@@ -337,7 +334,7 @@ jobs:
       matrix:
         os: [windows-latest]
         floatx: [float32]
-        python-version: ["3.12"]
+        python-version: ["3.13"]
         test-subset:
         - tests/sampling/test_mcmc.py tests/ode/test_ode.py tests/ode/test_utils.py tests/distributions/test_transform.py
       fail-fast: false
@@ -363,7 +360,6 @@ jobs:
       - name: Install-pymc
         run: |
           pip install -e .
-          pip install --pre -U 'polyagamma<1.3.7'
           python --version
           micromamba list
       - name: Run tests

conda-envs/environment-dev.yml

@@ -9,11 +9,10 @@ dependencies:
 - blas
 - cachetools>=4.2.1
 - cloudpickle
-- h5py>=2.7
 - numpy>=1.25.0
 - pandas>=0.24.0
 - pip
-- pytensor>=2.26.2,<2.28
+- pytensor>=2.28.1,<2.29
 - python-graphviz
 - networkx
 - scipy>=1.4.1
@@ -27,6 +26,7 @@ dependencies:
 - myst-nb<=1.0.0
 - numpydoc
 - pre-commit>=2.8.0
+- polyagamma
 - pytest-cov>=2.5
 - pytest>=3.0
 - rich>=13.7.1
@@ -36,9 +36,8 @@ dependencies:
 - sphinx>=1.5
 - sphinxext-rediraffe
 - watermark
-- polyagamma
 - sphinx-remove-toctrees
-- mypy=1.5.1
+- mypy=1.15.0
 - types-cachetools
 - pip:
   - git+https://github.com/pymc-devs/pymc-sphinx-theme

conda-envs/environment-docs.yml

@@ -11,7 +11,7 @@ dependencies:
 - numpy>=1.25.0
 - pandas>=0.24.0
 - pip
-- pytensor>=2.26.2,<2.28
+- pytensor>=2.28.1,<2.29
 - python-graphviz
 - rich>=13.7.1
 - scipy>=1.4.1

conda-envs/environment-jax.yml

@@ -9,7 +9,6 @@ dependencies:
 - blas
 - cachetools>=4.2.1
 - cloudpickle
-- h5py>=2.7
 - zarr>=2.5.0,<3
 # Jaxlib version must not be greater than jax version!
 - blackjax>=1.2.2
@@ -21,7 +20,7 @@ dependencies:
 - numpyro>=0.8.0
 - pandas>=0.24.0
 - pip
-- pytensor>=2.26.2,<2.28
+- pytensor>=2.28.1,<2.29
 - python-graphviz
 - networkx
 - rich>=13.7.1
@@ -34,7 +33,7 @@ dependencies:
 - pre-commit>=2.8.0
 - pytest-cov>=2.5
 - pytest>=3.0
-- mypy=1.5.1
+- mypy=1.15.0
 - types-cachetools
 - pip:
   - numdifftools>=0.9.40

conda-envs/environment-test.yml

@@ -9,12 +9,12 @@ dependencies:
 - blas
 - cachetools>=4.2.1
 - cloudpickle
-- h5py>=2.7
 - jax
 - numpy>=1.25.0
 - pandas>=0.24.0
 - pip
-- pytensor>=2.26.2,<2.28
+- polyagamma
+- pytensor>=2.28.1,<2.29
 - python-graphviz
 - networkx
 - rich>=13.7.1
@@ -27,7 +27,7 @@ dependencies:
 - pre-commit>=2.8.0
 - pytest-cov>=2.5
 - pytest>=3.0
-- mypy=1.5.1
+- mypy=1.15.0
 - types-cachetools
 - pip:
   - numdifftools>=0.9.40

conda-envs/windows-environment-dev.yml

@@ -9,11 +9,10 @@ dependencies:
 - blas
 - cachetools>=4.2.1
 - cloudpickle
-- h5py>=2.7
 - numpy>=1.25.0
 - pandas>=0.24.0
 - pip
-- pytensor>=2.26.2,<2.28
+- pytensor>=2.28.1,<2.29
 - python-graphviz
 - networkx
 - rich>=13.7.1
@@ -25,6 +24,7 @@ dependencies:
 - ipython>=7.16
 - myst-nb<=1.0.0
 - numpydoc
+- polyagamma
 - pre-commit>=2.8.0
 - pytest-cov>=2.5
 - pytest>=3.0
@@ -35,7 +35,7 @@ dependencies:
 - sphinx>=1.5
 - watermark
 - sphinx-remove-toctrees
-- mypy=1.5.1
+- mypy=1.15.0
 - types-cachetools
 - pip:
   - git+https://github.com/pymc-devs/pymc-sphinx-theme

conda-envs/windows-environment-test.yml

@@ -9,14 +9,13 @@ dependencies:
 - blas
 - cachetools>=4.2.1
 - cloudpickle
-- h5py>=2.7
 - libpython
 - mkl-service>=2.3.0
-- m2w64-toolchain
 - numpy>=1.25.0
 - pandas>=0.24.0
 - pip
-- pytensor>=2.26.2,<2.28
+- polyagamma
+- pytensor>=2.28.1,<2.29
 - python-graphviz
 - networkx
 - rich>=13.7.1
@@ -29,7 +28,7 @@ dependencies:
 - pre-commit>=2.8.0
 - pytest-cov>=2.5
 - pytest>=3.0
-- mypy=1.5.1
+- mypy=1.15.0
 - types-cachetools
 - pip:
   - numdifftools>=0.9.40

pymc/backends/mcbackend.py

@@ -181,7 +181,7 @@ def get_sampler_stats(
     def _slice(self, idx: slice) -> "IBaseTrace":
         # Get the integer indices
         start, stop, step = idx.indices(len(self))
-        indices = np.arange(start, stop, step)
+        indices = tuple(range(start, stop, step))
 
         # Create a NumPyChain for the sliced data
         nchain = mcb.backends.numpy.NumPyChain(

pymc/distributions/distribution.py

@@ -192,20 +192,19 @@ def support_point(op, rv, *dist_params):
         return new_cls
 
 
-class _class_or_instancemethod(classmethod):
-    """Allow a method to be called both as a classmethod and an instancemethod.
+class _class_or_instance_property(property):
+    """Allow a property to be accessed from a class or an instance.
 
-    Priority is given to the instancemethod.
+    Priority is given to the instance.
 
     This is used to allow extracting information from the signature of a SymbolicRandomVariable
-    which may be provided either as a class attribute or as an instance attribute.
+    which may be available early as a class attribute or only later as an instance attribute.
 
-    Adapted from https://stackoverflow.com/a/28238047
+    Adapted from https://stackoverflow.com/a/13624858
     """
 
-    def __get__(self, instance, type_):
-        descr_get = super().__get__ if instance is None else self.__func__.__get__
-        return descr_get(instance, type_)
+    def __get__(self, owner_self, owner_cls):
+        return self.fget(owner_self if owner_self is not None else owner_cls)
 
 
 class SymbolicRandomVariable(MeasurableOp, OpFromGraph):
@@ -241,8 +240,7 @@ class SymbolicRandomVariable(MeasurableOp, OpFromGraph):
     _print_name: tuple[str, str] = ("Unknown", "\\operatorname{Unknown}")
     """Tuple of (name, latex name) used for for pretty-printing variables of this type"""
 
-    @_class_or_instancemethod
-    @property
+    @_class_or_instance_property
     def signature(cls_or_self) -> None | str:
         # Convert "expanded" signature into "vanilla" signature that has no rng and size tokens
         extended_signature = cls_or_self.extended_signature
@@ -257,40 +255,28 @@ def signature(cls_or_self) -> None | str:
 
         return signature
 
-    @_class_or_instancemethod
-    @property
+    @_class_or_instance_property
     def ndims_params(cls_or_self) -> Sequence[int] | None:
-        """Number of core dimensions of the distribution's parameters."""
+        """Return number of core dimensions of the distribution's parameters."""
         signature = cls_or_self.signature
         if signature is None:
             return None
         inputs_signature, _ = _parse_gufunc_signature(signature)
         return [len(sig) for sig in inputs_signature]
 
-    @_class_or_instancemethod
-    @property
+    @_class_or_instance_property
     def ndim_supp(cls_or_self) -> int | None:
-        """Number of support dimensions of the RandomVariable.
+        """Return number of support dimensions of the RandomVariable.
 
         (0 for scalar, 1 for vector, ...)
         """
         signature = cls_or_self.signature
         if signature is None:
-            return None
+            return getattr(cls_or_self, "_ndim_supp", None)
         _, outputs_params_signature = _parse_gufunc_signature(signature)
         return max(len(out_sig) for out_sig in outputs_params_signature)
 
-    @_class_or_instancemethod
-    def _parse_extended_signature(cls_or_self) -> tuple[tuple[str, ...], tuple[str, ...]] | None:
-        extended_signature = cls_or_self.extended_signature
-        if extended_signature is None:
-            return None
-
-        fake_signature = extended_signature.replace("[rng]", "(rng)").replace("[size]", "(size)")
-        return _parse_gufunc_signature(fake_signature)
-
-    @_class_or_instancemethod
-    @property
+    @_class_or_instance_property
     def default_output(cls_or_self) -> int | None:
         extended_signature = cls_or_self.extended_signature
         if extended_signature is None:
@@ -374,7 +360,7 @@ def __init__(
         if "ndim_supp" in kwargs:
             # For backwards compatibility we allow passing ndim_supp without signature
             # This is the only variable that PyMC absolutely needs to work with SymbolicRandomVariables
-            self.ndim_supp = kwargs.pop("ndim_supp")
+            self._ndim_supp = kwargs.pop("ndim_supp")
 
         if self.ndim_supp is None:
             raise ValueError("ndim_supp or signature must be provided")

pymc/distributions/multivariate.py

@@ -2057,14 +2057,12 @@ class KroneckerNormal(Continuous):
     rv_op = KroneckerNormalRV.rv_op
 
     @classmethod
-    def dist(cls, mu, covs=None, chols=None, evds=None, sigma=None, *args, **kwargs):
+    def dist(cls, mu, covs=None, chols=None, evds=None, sigma=0.0, *args, **kwargs):
         if len([i for i in [covs, chols, evds] if i is not None]) != 1:
             raise ValueError(
                 "Incompatible parameterization. Specify exactly one of covs, chols, or evds."
             )
 
-        sigma = sigma if sigma else 0
-
         if chols is not None:
             covs = [chol.dot(chol.T) for chol in chols]
         elif evds is not None:
@@ -2076,6 +2074,7 @@ def dist(cls, mu, covs=None, chols=None, evds=None, sigma=None, *args, **kwargs)
                 covs.append(cov_i)
 
         mu = pt.as_tensor_variable(mu)
+        sigma = pt.as_tensor_variable(sigma)
 
         return super().dist([mu, sigma, *covs], **kwargs)
 

pymc/distributions/shape_utils.py

@@ -369,7 +369,7 @@ def get_support_shape(
         support_shape_offset = [0] * ndim_supp
     elif isinstance(support_shape_offset, int):
         support_shape_offset = [support_shape_offset] * ndim_supp
-    inferred_support_shape: Sequence[int | np.ndarray | Variable] | None = None
+    inferred_support_shape: Sequence[int | np.ndarray | TensorVariable] | None = None
 
     if shape is not None:
         shape = to_tuple(shape)
@@ -378,9 +378,7 @@ def get_support_shape(
             raise ValueError(
                 f"Number of shape dimensions is too small for ndim_supp of {ndim_supp}"
             )
-        inferred_support_shape = [
-            shape[i] - support_shape_offset[i] for i in np.arange(-ndim_supp, 0)
-        ]
+        inferred_support_shape = [shape[i] - support_shape_offset[i] for i in range(-ndim_supp, 0)]
 
     if inferred_support_shape is None and dims is not None:
         dims = convert_dims(dims)
@@ -389,7 +387,7 @@ def get_support_shape(
             raise ValueError(f"Number of dims is too small for ndim_supp of {ndim_supp}")
         model = modelcontext(None)
         inferred_support_shape = [
-            model.dim_lengths[dims[i]] - support_shape_offset[i] for i in np.arange(-ndim_supp, 0)
+            model.dim_lengths[dims[i]] - support_shape_offset[i] for i in range(-ndim_supp, 0)
         ]
 
     if inferred_support_shape is None and observed is not None:
@@ -399,7 +397,7 @@ def get_support_shape(
                 f"Number of observed dimensions is too small for ndim_supp of {ndim_supp}"
             )
         inferred_support_shape = [
-            observed.shape[i] - support_shape_offset[i] for i in np.arange(-ndim_supp, 0)
+            observed.shape[i] - support_shape_offset[i] for i in range(-ndim_supp, 0)
         ]
 
     if inferred_support_shape is None:
@@ -413,7 +411,7 @@ def get_support_shape(
         # There were two sources of support_shape, make sure they are consistent
         inferred_support_shape = [
             cast(
-                Variable,
+                TensorVariable,
                 Assert(msg="support_shape does not match respective shape dimension")(
                     inferred, pt.eq(inferred, explicit)
                 ),