diff --git a/.github/workflows/mypy.yml b/.github/workflows/mypy.yml index e6ea6826f8..4a2d072b9e 100644 --- a/.github/workflows/mypy.yml +++ b/.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: | diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 268656f68b..97e50bef2a 100644 --- a/.github/workflows/tests.yml +++ b/.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 diff --git a/conda-envs/environment-dev.yml b/conda-envs/environment-dev.yml index 546e1277d4..b2e8f17849 100644 --- a/conda-envs/environment-dev.yml +++ b/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 diff --git a/conda-envs/environment-docs.yml b/conda-envs/environment-docs.yml index fd035676c3..823f6563f6 100644 --- 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.26.2,<2.28 +- pytensor>=2.28.1,<2.29 - 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 6bde602133..ea630e806f 100644 --- a/conda-envs/environment-jax.yml +++ b/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 diff --git a/conda-envs/environment-test.yml b/conda-envs/environment-test.yml index 20f0478998..37729d4d1a 100644 --- a/conda-envs/environment-test.yml +++ b/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 diff --git a/conda-envs/windows-environment-dev.yml b/conda-envs/windows-environment-dev.yml index 503cf125b2..1c4d8f04d5 100644 --- a/conda-envs/windows-environment-dev.yml +++ b/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 diff --git a/conda-envs/windows-environment-test.yml b/conda-envs/windows-environment-test.yml index 5136d997a3..63f5f8ee4e 100644 --- a/conda-envs/windows-environment-test.yml +++ b/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 diff --git a/pymc/backends/mcbackend.py b/pymc/backends/mcbackend.py index 9bc8ff1043..d02a6dbebb 100644 --- a/pymc/backends/mcbackend.py +++ b/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( diff --git a/pymc/distributions/distribution.py b/pymc/distributions/distribution.py index b2ec6fb79b..86dc5ea0d4 100644 --- a/pymc/distributions/distribution.py +++ b/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") diff --git a/pymc/distributions/multivariate.py b/pymc/distributions/multivariate.py index 32f9e30f06..9ff56027a8 100644 --- a/pymc/distributions/multivariate.py +++ b/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) diff --git a/pymc/distributions/shape_utils.py b/pymc/distributions/shape_utils.py index 7dd0d94414..98c743b70e 100644 --- a/pymc/distributions/shape_utils.py +++ b/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) ), diff --git a/pymc/gp/cov.py b/pymc/gp/cov.py index 16e51cdde9..e275e8de52 100644 --- a/pymc/gp/cov.py +++ b/pymc/gp/cov.py @@ -164,12 +164,13 @@ class Covariance(BaseCovariance): def __init__(self, input_dim: int, active_dims: IntSequence | None = None): self.input_dim = input_dim + self.active_dims: np.ndarray[Any, np.dtype[np.int_]] if active_dims is None: self.active_dims = np.arange(input_dim) else: self.active_dims = np.asarray(active_dims, int) - if max(self.active_dims) > self.input_dim: + if self.active_dims.max() > self.input_dim: raise ValueError("Values in `active_dims` can't be larger than `input_dim`.") @property diff --git a/pymc/logprob/transforms.py b/pymc/logprob/transforms.py index 8446874996..8b5eac7b16 100644 --- a/pymc/logprob/transforms.py +++ b/pymc/logprob/transforms.py @@ -470,7 +470,9 @@ def find_measurable_transforms(fgraph: FunctionGraph, node: Apply) -> list[Varia # Do not apply rewrite to discrete variables except for their addition and negation if measurable_input.type.dtype.startswith("int"): - if not (find_negated_var(measurable_output) or isinstance(node.op.scalar_op, Add)): + if not ( + find_negated_var(measurable_output) is not None or isinstance(node.op.scalar_op, Add) + ): return None # Do not allow rewrite if output is cast to a float, because we don't have meta-info on the type of the MeasurableVariable if not measurable_output.type.dtype.startswith("int"): diff --git a/pymc/logprob/utils.py b/pymc/logprob/utils.py index e1fdc903ee..a028341032 100644 --- a/pymc/logprob/utils.py +++ b/pymc/logprob/utils.py @@ -186,13 +186,12 @@ def expand_fn(var): return [] if any( - ancestor_var - for ancestor_var in walk(inputs, expand=expand_fn, bfs=False) - if ( + ( ancestor_var.owner and isinstance(ancestor_var.owner.op, MeasurableOp) and not isinstance(ancestor_var.owner.op, ValuedRV) ) + for ancestor_var in walk(inputs, expand=expand_fn, bfs=False) ): return True return False diff --git a/pymc/model/core.py b/pymc/model/core.py index cef6cd6d1b..5a7b2651cf 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -895,7 +895,7 @@ def coords(self) -> dict[str, tuple | None]: return self._coords @property - def dim_lengths(self) -> dict[str, Variable]: + def dim_lengths(self) -> dict[str, TensorVariable]: """The symbolic lengths of dimensions in the model. The values are typically instances of ``TensorVariable`` or ``ScalarSharedVariable``. diff --git a/pymc/sampling/forward.py b/pymc/sampling/forward.py index 19957e0540..5f3c983b58 100644 --- a/pymc/sampling/forward.py +++ b/pymc/sampling/forward.py @@ -876,7 +876,7 @@ def sample_posterior_predictive( try: with progress: task = progress.add_task("Sampling ...", completed=0, total=samples) - for idx in np.arange(samples): + for idx in range(samples): if nchain > 1: # the trace object will either be a MultiTrace (and have _straces)... if hasattr(_trace, "_straces"): diff --git a/requirements-dev.txt b/requirements-dev.txt index b0ef92f505..c868dbd52e 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -5,11 +5,10 @@ arviz>=0.13.0 cachetools>=4.2.1 cloudpickle git+https://github.com/pymc-devs/pymc-sphinx-theme -h5py>=2.7 ipython>=7.16 jupyter-sphinx mcbackend>=0.4.0 -mypy==1.5.1 +mypy==1.15.0 myst-nb<=1.0.0 numdifftools>=0.9.40 numpy>=1.25.0 @@ -17,7 +16,7 @@ numpydoc pandas>=0.24.0 polyagamma pre-commit>=2.8.0 -pytensor>=2.26.2,<2.28 +pytensor>=2.28.1,<2.29 pytest-cov>=2.5 pytest>=3.0 rich>=13.7.1 diff --git a/requirements.txt b/requirements.txt index c1f82bb4dd..fb0c131ce7 100644 --- 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.26.1,<2.28 +pytensor>=2.28.1,<2.29 rich>=13.7.1 scipy>=1.4.1 threadpoolctl>=3.1.0,<4.0.0 diff --git a/setup.py b/setup.py index 99bcadd86a..30e50981c3 100755 --- a/setup.py +++ b/setup.py @@ -33,6 +33,7 @@ "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", "License :: OSI Approved :: Apache Software License", "Intended Audience :: Science/Research", "Topic :: Scientific/Engineering", diff --git a/tests/distributions/test_continuous.py b/tests/distributions/test_continuous.py index cfdd0b3d60..9c108d2035 100644 --- a/tests/distributions/test_continuous.py +++ b/tests/distributions/test_continuous.py @@ -2338,10 +2338,8 @@ class TestInverseGamma(BaseTestDistributionRandom): pymc_dist_params = {"alpha": 2.0, "beta": 5.0} expected_rv_op_params = {"alpha": 2.0, "beta": 5.0} reference_dist_params = {"a": 2.0, "scale": 5.0} - reference_dist = seeded_scipy_distribution_builder("invgamma") checks_to_run = [ "check_pymc_params_match_rv_op", - "check_pymc_draws_match_reference", ] @@ -2440,7 +2438,9 @@ def test_rng_different_shapes(self): ) class TestPolyaGamma(BaseTestDistributionRandom): def polyagamma_rng_fn(self, size, h, z, rng): - return random_polyagamma(h, z, size=size, random_state=rng._bit_generator) + # Polyagamma returns different values if inputs have explicit broadcasted dims + # Which PyTensor RVs always do when size is not None. + return random_polyagamma(np.atleast_1d(h), np.atleast_1d(z), size=size, random_state=rng) pymc_dist = pm.PolyaGamma pymc_dist_params = {"h": 1.0, "z": 0.0} diff --git a/tests/distributions/test_custom.py b/tests/distributions/test_custom.py index d3de7cf4f7..6546aa22e5 100644 --- a/tests/distributions/test_custom.py +++ b/tests/distributions/test_custom.py @@ -603,7 +603,7 @@ def test_symbolic_dist(self): def dist(size): return Truncated.dist(Beta.dist(1, 1, size=size), lower=0.1, upper=0.9) - assert CustomDist.dist(dist=dist) + CustomDist.dist(dist=dist) def test_nested_custom_dist(self): """Test we can create CustomDist that creates another CustomDist""" diff --git a/tests/distributions/test_multivariate.py b/tests/distributions/test_multivariate.py index 39b6c562e1..2694230c32 100644 --- a/tests/distributions/test_multivariate.py +++ b/tests/distributions/test_multivariate.py @@ -27,6 +27,7 @@ from pytensor.tensor import TensorVariable from pytensor.tensor.blockwise import Blockwise from pytensor.tensor.nlinalg import MatrixInverse +from pytensor.tensor.random.basic import multivariate_normal from pytensor.tensor.random.utils import broadcast_params from pytensor.tensor.slinalg import Cholesky @@ -421,7 +422,7 @@ def test_matrixnormal(self, n): @pytest.mark.parametrize("n", [2, 3]) @pytest.mark.parametrize("m", [3]) - @pytest.mark.parametrize("sigma", [None, 1]) + @pytest.mark.parametrize("sigma", [0, 1]) def test_kroneckernormal(self, n, m, sigma): np.random.seed(5) N = n * m @@ -1392,6 +1393,11 @@ def test_dirichlet_multinomial_support_point(self, a, n, size, expected): class TestMvNormalCov(BaseTestDistributionRandom): + def mvnormal_rng_fn(self, size, mean, cov, rng): + if isinstance(size, int): + size = (size,) + return multivariate_normal.rng_fn(rng, mean, cov, size=size) + pymc_dist = pm.MvNormal pymc_dist_params = { "mu": np.array([1.0, 2.0]), @@ -1407,7 +1413,8 @@ class TestMvNormalCov(BaseTestDistributionRandom): "mean": np.array([1.0, 2.0]), "cov": np.array([[2.0, 0.0], [0.0, 3.5]]), } - reference_dist = seeded_numpy_distribution_builder("multivariate_normal") + reference_dist = lambda self: ft.partial(self.mvnormal_rng_fn, rng=self.get_random_state()) # noqa: E731 + checks_to_run = [ "check_pymc_params_match_rv_op", "check_pymc_draws_match_reference", @@ -1531,12 +1538,12 @@ class TestZeroSumNormal: def assert_zerosum_axes(self, random_samples, axes_to_check, check_zerosum_axes=True): if check_zerosum_axes: for ax in axes_to_check: - assert np.isclose(random_samples.mean(axis=ax), 0).all(), ( + assert np.allclose(random_samples.mean(axis=ax), 0), ( f"{ax} is a zerosum_axis but is not summing to 0 across all samples." ) else: for ax in axes_to_check: - assert not np.isclose(random_samples.mean(axis=ax), 0).all(), ( + assert not np.allclose(random_samples.mean(axis=ax), 0), ( f"{ax} is not a zerosum_axis, but is nonetheless summing to 0 across all samples." ) @@ -1564,8 +1571,8 @@ def test_zsn_dims(self, dims, n_zerosum_axes): ) ndim_supp = v.owner.op.ndim_supp - n_zerosum_axes = np.arange(-ndim_supp, 0) - nonzero_axes = np.arange(v.ndim - ndim_supp) + n_zerosum_axes = tuple(range(-ndim_supp, 0)) + nonzero_axes = tuple(range(v.ndim - ndim_supp)) for samples in [ s.posterior.v, random_samples, @@ -1595,8 +1602,8 @@ def test_zsn_shape(self, n_zerosum_axes): ) ndim_supp = v.owner.op.ndim_supp - n_zerosum_axes = np.arange(-ndim_supp, 0) - nonzero_axes = np.arange(v.ndim - ndim_supp) + n_zerosum_axes = tuple(range(-ndim_supp, 0)) + nonzero_axes = tuple(range(v.ndim - ndim_supp)) for samples in [ s.posterior.v, random_samples, @@ -1775,7 +1782,9 @@ def test_batched_sigma(self): class TestMvStudentTCov(BaseTestDistributionRandom): def mvstudentt_rng_fn(self, size, nu, mu, scale, rng): - mv_samples = rng.multivariate_normal(np.zeros_like(mu), scale, size=size) + if isinstance(size, int): + size = (size,) + mv_samples = multivariate_normal.rng_fn(rng, np.zeros_like(mu), scale, size=size) chi2_samples = rng.chisquare(nu, size=size) return (mv_samples / np.sqrt(chi2_samples[:, None] / nu)) + mu @@ -2111,9 +2120,11 @@ def check_random_variable_prior(self): class TestKroneckerNormal(BaseTestDistributionRandom): def kronecker_rng_fn(self, size, mu, covs=None, sigma=None, rng=None): - cov = pm.math.kronecker(covs[0], covs[1]).eval() + if isinstance(size, int): + size = (size,) + cov = np.kron(covs[0], covs[1]) cov += sigma**2 * np.identity(cov.shape[0]) - return st.multivariate_normal.rvs(mean=mu, cov=cov, size=size, random_state=rng) + return multivariate_normal.rng_fn(rng, mean=mu, cov=cov, size=size) pymc_dist = pm.KroneckerNormal diff --git a/tests/logprob/test_basic.py b/tests/logprob/test_basic.py index 64cbf63b3e..f6014d78a3 100644 --- a/tests/logprob/test_basic.py +++ b/tests/logprob/test_basic.py @@ -382,7 +382,7 @@ def test_warn_random_found_probability_inference(func, scipy_func, test_value): with pytest.warns( UserWarning, match="RandomVariables {input} were found in the derived graph" ): - assert func(rv, 0.0) + func(rv, 0.0) res = func(rv, 0.0, warn_rvs=False) # This is the problem we are warning about, as now we can no longer identify the original rv in the graph diff --git a/tests/logprob/test_order.py b/tests/logprob/test_order.py index 1b1fbac636..32c5d3c001 100644 --- a/tests/logprob/test_order.py +++ b/tests/logprob/test_order.py @@ -291,4 +291,4 @@ def test_non_measurable_max_grad(): joint_logp = pt.sum([term.sum() for term in logp_terms]) # Test that calling gradient does not raise a NotImplementedError - assert pt.grad(joint_logp, x_vv) + pt.grad(joint_logp, x_vv)