Bump PyTensor dependency and support Python 3.14

.github/workflows/tests.yml

@@ -58,7 +58,7 @@ jobs:
       matrix:
         os: [ubuntu-latest]
         floatx: [float64]
-        python-version: ["3.13"]
+        python-version: ["3.14"]
         test-subset:
           - |
             tests/test_util.py
@@ -237,7 +237,7 @@ jobs:
       matrix:
         os: [macos-latest]
         floatx: [float64]
-        python-version: ["3.13"]
+        python-version: ["3.14"]
         test-subset:
           - |
             tests/sampling/test_parallel.py
@@ -295,6 +295,8 @@ jobs:
       matrix:
         os: [ubuntu-latest]
         floatx: [float64]
+        # nutpie depends on PyMC, and it will require an extra release cycle to support
+        # the next PyMC release and therefore Python 3.14.
         python-version: ["3.13"]
         test-subset:
           - |
@@ -345,7 +347,7 @@ jobs:
       matrix:
         os: [windows-latest]
         floatx: [float32]
-        python-version: ["3.13"]
+        python-version: ["3.14"]
         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

conda-envs/environment-alternative-backends.yml

@@ -22,7 +22,7 @@ dependencies:
 - numpyro>=0.8.0
 - pandas>=0.24.0
 - pip
-- pytensor>=2.35.0,<2.36
+- pytensor>=2.36.0,<2.37
 - python-graphviz
 - networkx
 - rich>=13.7.1
@@ -35,7 +35,7 @@ dependencies:
 - pre-commit>=2.8.0
 - pytest-cov>=2.5
 - pytest>=3.0
-- mypy=1.15.0
+- mypy=1.19.1
 - types-cachetools
 - pip:
   - numdifftools>=0.9.40

conda-envs/environment-dev.yml

@@ -12,7 +12,7 @@ dependencies:
 - numpy>=1.25.0
 - pandas>=0.24.0
 - pip
-- pytensor>=2.35.0,<2.36
+- pytensor>=2.36.0,<2.37
 - python-graphviz
 - networkx
 - scipy>=1.4.1
@@ -37,7 +37,7 @@ dependencies:
 - sphinxext-rediraffe
 - watermark
 - sphinx-remove-toctrees
-- mypy=1.15.0
+- mypy=1.19.1
 - types-cachetools
 - pip:
   - pymc-sphinx-theme>=0.16.0

conda-envs/environment-docs.yml

@@ -11,7 +11,7 @@ dependencies:
 - numpy>=1.25.0
 - pandas>=0.24.0
 - pip
-- pytensor>=2.35.0,<2.36
+- pytensor>=2.36.0,<2.37
 - 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.35.0,<2.36
+- pytensor>=2.36.0,<2.37
 - 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.15.0
+- mypy=1.19.1
 - types-cachetools
 - pip:
   - numdifftools>=0.9.40

conda-envs/windows-environment-dev.yml

@@ -12,7 +12,7 @@ dependencies:
 - numpy>=1.25.0
 - pandas>=0.24.0
 - pip
-- pytensor>=2.35.0,<2.36
+- pytensor>=2.36.0,<2.37
 - python-graphviz
 - networkx
 - rich>=13.7.1
@@ -35,7 +35,7 @@ dependencies:
 - sphinx>=1.5
 - watermark
 - sphinx-remove-toctrees
-- mypy=1.15.0
+- mypy=1.19.1
 - types-cachetools
 - pip:
   - git+https://github.com/pymc-devs/pymc-sphinx-theme

conda-envs/windows-environment-test.yml

@@ -15,7 +15,7 @@ dependencies:
 - pandas>=0.24.0
 - pip
 - polyagamma
-- pytensor>=2.35.0,<2.36
+- pytensor>=2.36.0,<2.37
 - python-graphviz
 - networkx
 - rich>=13.7.1
@@ -28,7 +28,7 @@ dependencies:
 - pre-commit>=2.8.0
 - pytest-cov>=2.5
 - pytest>=3.0
-- mypy=1.15.0
+- mypy=1.19.1
 - types-cachetools
 - pip:
   - numdifftools>=0.9.40

pymc/distributions/multivariate.py

@@ -25,7 +25,8 @@
 from pytensor.graph.basic import Apply, Variable
 from pytensor.graph.op import Op
 from pytensor.raise_op import Assert
-from pytensor.sparse.basic import DenseFromSparse, sp_sum
+from pytensor.sparse.basic import DenseFromSparse
+from pytensor.sparse.math import sp_sum
 from pytensor.tensor import (
     TensorConstant,
     TensorVariable,
@@ -2263,10 +2264,12 @@ class CAR(Continuous):
     def dist(cls, mu, W, alpha, tau, *args, **kwargs):
         # This variable has an expensive validation check, that we want to constant-fold if possible
         # So it's passed as an explicit input
-        W = pytensor.sparse.as_sparse_or_tensor_variable(W)
+        from pytensor.sparse import as_sparse_or_tensor_variable, structured_sign
+
+        W = as_sparse_or_tensor_variable(W)
         if isinstance(W.type, pytensor.sparse.SparseTensorType):
-            abs_diff = pytensor.sparse.basic.mul(pytensor.sparse.sign(W - W.T), W - W.T)
-            W_is_valid = pt.isclose(pytensor.sparse.sp_sum(abs_diff), 0)
+            abs_diff = structured_sign(W - W.T) * (W - W.T)
+            W_is_valid = pt.isclose(abs_diff.sum(), 0)
         else:
             W_is_valid = pt.allclose(W, W.T)
 
@@ -2307,7 +2310,7 @@ def logp(value, mu, W, alpha, tau, W_is_valid):
         if W.owner and isinstance(W.owner.op, DenseFromSparse):
             W = W.owner.inputs[0]
 
-        sparse = isinstance(W, pytensor.sparse.SparseVariable)
+        sparse = isinstance(W, pytensor.sparse.variable.SparseVariable)
         if sparse:
             D = sp_sum(W, axis=0)
             Dinv_sqrt = pt.diag(1 / pt.sqrt(D))

pymc/distributions/shape_utils.py

@@ -409,7 +409,7 @@ def get_support_shape(
         ]
 
     if inferred_support_shape is None and observed is not None:
-        observed = convert_observed_data(observed)
+        observed = cast(TensorVariable | np.ndarray, convert_observed_data(observed))
         if observed.ndim < ndim_supp:
             raise ValueError(
                 f"Number of observed dimensions is too small for ndim_supp of {ndim_supp}"

pymc/distributions/simulator.py

@@ -338,7 +338,7 @@ def make_node(self, x):
 
         def perform(self, node, inputs, outputs):
             (x,) = inputs
-            outputs[0][0] = np.atleast_1d(fn(x)).astype(pytensor.config.floatX)
+            outputs[0][0] = np.atleast_1d(fn(x)).astype(node.outputs[0].dtype)
 
     return SumStat()
 
@@ -365,8 +365,6 @@ def make_node(self, epsilon, obs_data, sim_data):
 
         def perform(self, node, inputs, outputs):
             eps, obs_data, sim_data = inputs
-            outputs[0][0] = np.atleast_1d(fn(eps, obs_data, sim_data)).astype(
-                pytensor.config.floatX
-            )
+            outputs[0][0] = np.atleast_1d(fn(eps, obs_data, sim_data)).astype(node.outputs[0].dtype)
 
     return Distance()

pymc/distributions/timeseries.py

@@ -442,23 +442,26 @@ def rv_op(cls, rhos, sigma, init_dist, steps, ar_order, constant_term, size=None
         rhos_bcast = pt.broadcast_to(rhos, rhos_bcast_shape)
 
         def step(*args):
-            *prev_xs, reversed_rhos, sigma, rng = args
+            *prev_xs, rng, reversed_rhos, sigma = args
             if constant_term:
                 mu = reversed_rhos[-1] + pt.sum(prev_xs * reversed_rhos[:-1], axis=0)
             else:
                 mu = pt.sum(prev_xs * reversed_rhos, axis=0)
             next_rng, new_x = Normal.dist(mu=mu, sigma=sigma, rng=rng).owner.outputs
-            return new_x, {rng: next_rng}
+            return new_x, next_rng
 
         # We transpose inputs as scan iterates over first dimension
-        innov, innov_updates = pytensor.scan(
+        innov, noise_next_rng = pytensor.scan(
             fn=step,
-            outputs_info=[{"initial": init_dist.T, "taps": range(-ar_order, 0)}],
-            non_sequences=[rhos_bcast.T[::-1], sigma.T, noise_rng],
+            outputs_info=[
+                {"initial": init_dist.T, "taps": range(-ar_order, 0)},
+                noise_rng,
+            ],
+            non_sequences=[rhos_bcast.T[::-1], sigma.T],
             n_steps=steps,
             strict=True,
+            return_updates=False,
         )
-        (noise_next_rng,) = tuple(innov_updates.values())
         ar = pt.concatenate([init_dist, innov.T], axis=-1)
 
         return AutoRegressiveRV(
@@ -710,24 +713,25 @@ def rv_op(cls, omega, alpha_1, beta_1, initial_vol, init_dist, steps, size=None)
 
         # Create OpFromGraph representing random draws from GARCH11 process
 
-        def step(prev_y, prev_sigma, omega, alpha_1, beta_1, rng):
+        def step(prev_y, prev_sigma, rng, omega, alpha_1, beta_1):
             new_sigma = pt.sqrt(
                 omega + alpha_1 * pt.square(prev_y) + beta_1 * pt.square(prev_sigma)
             )
             next_rng, new_y = Normal.dist(mu=0, sigma=new_sigma, rng=rng).owner.outputs
-            return (new_y, new_sigma), {rng: next_rng}
+            return new_y, new_sigma, next_rng
 
-        (y_t, _), innov_updates = pytensor.scan(
+        y_t, _, noise_next_rng = pytensor.scan(
             fn=step,
             outputs_info=[
                 init_dist,
                 pt.broadcast_to(initial_vol.astype("floatX"), init_dist.shape),
+                noise_rng,
             ],
-            non_sequences=[omega, alpha_1, beta_1, noise_rng],
+            non_sequences=[omega, alpha_1, beta_1],
             n_steps=steps,
             strict=True,
+            return_updates=False,
         )
-        (noise_next_rng,) = tuple(innov_updates.values())
 
         garch11 = pt.concatenate([init_dist[None, ...], y_t], axis=0).dimshuffle(
             (*range(1, y_t.ndim), 0)
@@ -816,12 +820,13 @@ def garch11_logp(
     def volatility_update(x, vol, w, a, b):
         return pt.sqrt(w + a * pt.square(x) + b * pt.square(vol))
 
-    vol, _ = pytensor.scan(
+    vol = pytensor.scan(
         fn=volatility_update,
         sequences=[value_dimswapped[:-1]],
         outputs_info=[initial_vol],
         non_sequences=[omega, alpha_1, beta_1],
         strict=True,
+        return_updates=False,
     )
     sigma_t = pt.concatenate([[initial_vol], vol])
     # Compute and collapse logp across time dimension
@@ -861,21 +866,21 @@ def rv_op(cls, init_dist, steps, sde_pars, dt, sde_fn, size=None):
 
         # Create OpFromGraph representing random draws from SDE process
         def step(*prev_args):
-            prev_y, *prev_sde_pars, rng = prev_args
+            prev_y, rng, *prev_sde_pars = prev_args
             f, g = sde_fn(prev_y, *prev_sde_pars)
             mu = prev_y + dt * f
             sigma = pt.sqrt(dt) * g
             next_rng, next_y = Normal.dist(mu=mu, sigma=sigma, rng=rng).owner.outputs
-            return next_y, {rng: next_rng}
+            return next_y, next_rng
 
-        y_t, innov_updates = pytensor.scan(
+        y_t, noise_next_rng = pytensor.scan(
             fn=step,
-            outputs_info=[init_dist],
-            non_sequences=[*sde_pars, noise_rng],
+            outputs_info=[init_dist, noise_rng],
+            non_sequences=[*sde_pars],
             n_steps=steps,
             strict=True,
+            return_updates=False,
         )
-        (noise_next_rng,) = tuple(innov_updates.values())
 
         sde_out = pt.concatenate([init_dist[None, ...], y_t], axis=0).dimshuffle(
             (*range(1, y_t.ndim), 0)

pymc/logprob/transforms.py

@@ -729,11 +729,12 @@ def calc_delta_x(value, prior_result):
                 2 * prior_result * pt.erfcx(prior_result) - 2 / pt.sqrt(np.pi)
             )
 
-        result, updates = scan(
+        result = scan(
             fn=calc_delta_x,
             outputs_info=pt.ones_like(x),
             non_sequences=value,
             n_steps=10,
+            return_updates=False,
         )
         return result[-1]
 

pymc/pytensorf.py

@@ -331,7 +331,7 @@ def jacobian1(f, v):
     def grad_i(i):
         return gradient1(f[i], v)
 
-    return pytensor.map(grad_i, idx)[0]
+    return pytensor.map(grad_i, idx, return_updates=False)
 
 
 def jacobian(f, vars=None):
@@ -355,8 +355,13 @@ def grad_ii(i, f, x):
         return grad(f[i], x)[i]
 
     return pytensor.scan(
-        grad_ii, sequences=[idx], n_steps=f.shape[0], non_sequences=[f, x], name="jacobian_diag"
-    )[0]
+        grad_ii,
+        sequences=[idx],
+        n_steps=f.shape[0],
+        non_sequences=[f, x],
+        name="jacobian_diag",
+        return_updates=False,
+    )
 
 
 @pytensor.config.change_flags(compute_test_value="ignore")
@@ -381,7 +386,7 @@ def hessian_diag1(f, v):
     def hess_ii(i):
         return gradient1(g[i], v)[i]
 
-    return pytensor.map(hess_ii, idx)[0]
+    return pytensor.map(hess_ii, idx, return_updates=False)
 
 
 @pytensor.config.change_flags(compute_test_value="ignore")

pymc/smc/sampling.py

@@ -274,7 +274,7 @@ def _save_sample_stats(
         if idata_kwargs is not None:
             ikwargs.update(idata_kwargs)
         idata = to_inference_data(trace, **ikwargs)
-        idata = InferenceData(**idata, sample_stats=sample_stats)
+        idata = InferenceData(**idata, sample_stats=sample_stats)  # type: ignore[arg-type]
 
     return sample_stats, idata
 

pymc/variational/approximations.py

@@ -392,7 +392,10 @@ def evaluate_over_trace(self, node):
         def sample(post, *_):
             return graph_replace(node, {self.input: post}, strict=False)
 
-        nodes, _ = pytensor.scan(
-            sample, self.histogram, non_sequences=_known_scan_ignored_inputs(makeiter(node))
+        nodes = pytensor.scan(
+            sample,
+            self.histogram,
+            non_sequences=_known_scan_ignored_inputs(makeiter(node)),
+            return_updates=False,
         )
         return nodes

pymc/variational/operators.py

@@ -134,7 +134,7 @@ class KSD(Operator):
     has_test_function = True
     returns_loss = False
     require_logq = False
-    objective_class = KSDObjective
+    objective_class = KSDObjective  # type: ignore[assignment]
 
     def __init__(self, approx, temperature=1):
         super().__init__(approx)