pytensor 2.28.3__tar.gz → 2.30.0__tar.gz

{pytensor-2.28.3/pytensor.egg-info → pytensor-2.30.0}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: pytensor
-Version: 2.28.3
+Version: 2.30.0
 Summary: Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs.
 Author-email: pymc-devs <pymc.devs@gmail.com>
 License: .. _license:
@@ -108,6 +108,7 @@ Requires-Dist: jaxlib; extra == "jax"
 Provides-Extra: numba
 Requires-Dist: numba>=0.57; extra == "numba"
 Requires-Dist: llvmlite; extra == "numba"
+Dynamic: license-file
 
 .. image:: https://cdn.rawgit.com/pymc-devs/pytensor/main/doc/images/PyTensor_RGB.svg
     :height: 100px

{pytensor-2.28.3 → pytensor-2.30.0}/doc/library/tensor/basic.rst

@@ -619,9 +619,8 @@ dimensions, see :meth:`_tensor_py_operators.dimshuffle`.
 
 .. function:: shape_padleft(x, n_ones=1)
 
-    Reshape `x` by left padding the shape with `n_ones` 1s. Note that all
-    this new dimension will be broadcastable. To make them non-broadcastable
-    see the :func:`unbroadcast`.
+    Reshape `x` by left padding the shape with `n_ones` 1s.
+    All new dimensions will be broadcastable.
 
     :param x: variable to be reshaped
     :type x: any `TensorVariable` (or compatible)
@@ -633,9 +632,8 @@ dimensions, see :meth:`_tensor_py_operators.dimshuffle`.
 
 .. function:: shape_padright(x, n_ones=1)
 
-    Reshape `x` by right padding the shape with `n_ones` ones. Note that all
-    this new dimension will be broadcastable. To make them non-broadcastable
-    see the :func:`unbroadcast`.
+    Reshape `x` by right padding the shape with `n_ones` ones.
+    All new dimensions will be broadcastable.
 
     :param x: variable to be reshaped
     :type x: any TensorVariable (or compatible)
@@ -646,9 +644,8 @@ dimensions, see :meth:`_tensor_py_operators.dimshuffle`.
 
 .. function:: shape_padaxis(t, axis)
 
-    Reshape `t` by inserting ``1`` at the dimension `axis`. Note that this new
-    dimension will be broadcastable. To make it non-broadcastable
-    see the :func:`unbroadcast`.
+    Reshape `t` by inserting ``1`` at the dimension `axis`.
+    All new dimensions will be broadcastable.
 
     :type x: any `TensorVariable` (or compatible)
     :param x: variable to be reshaped

{pytensor-2.28.3 → pytensor-2.30.0}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2025-03-05T13:02:56+0100",
+ "date": "2025-03-28T02:16:52+0800",
  "dirty": false,
  "error": null,
- "full-revisionid": "89d5366c888f430003aa6f16a13ad9cdd2826d95",
- "version": "2.28.3"
+ "full-revisionid": "0b56ed9c6514e4ba641cb07fd7bbca629dd797c3",
+ "version": "2.30.0"
 }
 '''  # END VERSION_JSON
 

{pytensor-2.28.3 → pytensor-2.30.0}/pytensor/compile/function/pfunc.py

@@ -292,14 +292,8 @@ def rebuild_collect_shared(
                 f" shared_var.type={store_into.type},"
                 f" update_val={update_val}, update_val.type={getattr(update_val, 'type', None)})."
             )
-            err_sug = (
-                "If the difference is related to the broadcast pattern,"
-                " you can call the"
-                " tensor.shape.unbroadcast(var, axis_to_unbroadcast[, ...])"
-                " function to mask broadcastable dimensions."
-            )
 
-            raise TypeError(err_msg, err_sug)
+            raise TypeError(err_msg)
         assert store_into.type.is_super(update_val.type)
 
         update_d[store_into] = update_val

{pytensor-2.28.3 → pytensor-2.30.0}/pytensor/compile/mode.py

@@ -454,6 +454,19 @@ else:
         RewriteDatabaseQuery(include=["fast_run", "py_only"]),
     )
 
+NUMBA = Mode(
+    NumbaLinker(),
+    RewriteDatabaseQuery(
+        include=["fast_run", "numba"],
+        exclude=[
+            "cxx_only",
+            "BlasOpt",
+            "local_careduce_fusion",
+            "scan_save_mem_prealloc",
+        ],
+    ),
+)
+
 JAX = Mode(
     JAXLinker(),
     RewriteDatabaseQuery(
@@ -463,6 +476,7 @@ JAX = Mode(
             "BlasOpt",
             "fusion",
             "inplace",
+            "scan_save_mem_prealloc",
         ],
     ),
 )
@@ -476,16 +490,10 @@ PYTORCH = Mode(
             "fusion",
             "inplace",
             "local_uint_constant_indices",
+            "scan_save_mem_prealloc",
         ],
     ),
 )
-NUMBA = Mode(
-    NumbaLinker(),
-    RewriteDatabaseQuery(
-        include=["fast_run", "numba"],
-        exclude=["cxx_only", "BlasOpt", "local_careduce_fusion"],
-    ),
-)
 
 
 predefined_modes = {

{pytensor-2.28.3 → pytensor-2.30.0}/pytensor/configdefaults.py

@@ -1085,7 +1085,9 @@ def add_scan_configvars():
         "scan__allow_output_prealloc",
         "Allow/disallow memory preallocation for outputs inside of scan "
         "(default: True)",
-        BoolParam(True),
+        # Non-mutable because ScanSaveMem rewrite checks it,
+        # and we can't have the rewrite and the implementation mismatch
+        BoolParam(True, mutable=False),
         in_c_key=False,
     )
 

{pytensor-2.28.3 → pytensor-2.30.0}/pytensor/graph/op.py

@@ -366,7 +366,7 @@ class Op(MetaObject):
         .. [1] Giles, Mike. 2008. “An Extended Collection of Matrix Derivative Results for Forward and Reverse Mode Automatic Differentiation.”
 
         """
-        raise NotImplementedError()
+        raise NotImplementedError(f"grad not implemented for Op {self}")
 
     def L_op(
         self,

{pytensor-2.28.3 → pytensor-2.30.0}/pytensor/ifelse.py

@@ -26,7 +26,7 @@ from pytensor.graph.op import _NoPythonOp
 from pytensor.graph.replace import clone_replace
 from pytensor.graph.rewriting.basic import GraphRewriter, in2out, node_rewriter
 from pytensor.graph.type import HasDataType, HasShape
-from pytensor.tensor.shape import Reshape, Shape, SpecifyShape, Unbroadcast
+from pytensor.tensor.shape import Reshape, Shape, SpecifyShape
 
 
 if TYPE_CHECKING:
@@ -481,7 +481,6 @@ acceptable_ops = (
     Shape,
     SpecifyShape,
     Reshape,
-    Unbroadcast,
     pt.math.Dot,
     pt.math.Max,
     pt.math.Argmax,

{pytensor-2.28.3 → pytensor-2.30.0}/pytensor/link/jax/dispatch/__init__.py

@@ -14,6 +14,7 @@ import pytensor.link.jax.dispatch.random
 import pytensor.link.jax.dispatch.scalar
 import pytensor.link.jax.dispatch.scan
 import pytensor.link.jax.dispatch.shape
+import pytensor.link.jax.dispatch.signal
 import pytensor.link.jax.dispatch.slinalg
 import pytensor.link.jax.dispatch.sort
 import pytensor.link.jax.dispatch.sparse

{pytensor-2.28.3 → pytensor-2.30.0}/pytensor/link/jax/dispatch/random.py

@@ -1,6 +1,7 @@
 from functools import singledispatch
 
 import jax
+import jax.numpy as jnp
 import numpy as np
 from numpy.random import Generator
 from numpy.random.bit_generator import (  # type: ignore[attr-defined]
@@ -105,14 +106,24 @@ def jax_funcify_RandomVariable(op: ptr.RandomVariable, node, **kwargs):
         assert_size_argument_jax_compatible(node)
 
         def sample_fn(rng, size, *parameters):
-            return jax_sample_fn(op, node=node)(rng, size, out_dtype, *parameters)
+            rng_key = rng["jax_state"]
+            rng_key, sampling_key = jax.random.split(rng_key, 2)
+            rng["jax_state"] = rng_key
+            sample = jax_sample_fn(op, node=node)(
+                sampling_key, size, out_dtype, *parameters
+            )
+            return (rng, sample)
 
     else:
 
         def sample_fn(rng, size, *parameters):
-            return jax_sample_fn(op, node=node)(
-                rng, static_size, out_dtype, *parameters
+            rng_key = rng["jax_state"]
+            rng_key, sampling_key = jax.random.split(rng_key, 2)
+            rng["jax_state"] = rng_key
+            sample = jax_sample_fn(op, node=node)(
+                sampling_key, static_size, out_dtype, *parameters
             )
+            return (rng, sample)
 
     return sample_fn
 
@@ -133,12 +144,9 @@ def jax_sample_fn_generic(op, node):
     name = op.name
     jax_op = getattr(jax.random, name)
 
-    def sample_fn(rng, size, dtype, *parameters):
-        rng_key = rng["jax_state"]
-        rng_key, sampling_key = jax.random.split(rng_key, 2)
-        sample = jax_op(sampling_key, *parameters, shape=size, dtype=dtype)
-        rng["jax_state"] = rng_key
-        return (rng, sample)
+    def sample_fn(rng_key, size, dtype, *parameters):
+        sample = jax_op(rng_key, *parameters, shape=size, dtype=dtype)
+        return sample
 
     return sample_fn
 
@@ -159,29 +167,23 @@ def jax_sample_fn_loc_scale(op, node):
     name = op.name
     jax_op = getattr(jax.random, name)
 
-    def sample_fn(rng, size, dtype, *parameters):
-        rng_key = rng["jax_state"]
-        rng_key, sampling_key = jax.random.split(rng_key, 2)
+    def sample_fn(rng_key, size, dtype, *parameters):
         loc, scale = parameters
         if size is None:
             size = jax.numpy.broadcast_arrays(loc, scale)[0].shape
-        sample = loc + jax_op(sampling_key, size, dtype) * scale
-        rng["jax_state"] = rng_key
-        return (rng, sample)
+        sample = loc + jax_op(rng_key, size, dtype) * scale
+        return sample
 
     return sample_fn
 
 
 @jax_sample_fn.register(ptr.MvNormalRV)
 def jax_sample_mvnormal(op, node):
-    def sample_fn(rng, size, dtype, mean, cov):
-        rng_key = rng["jax_state"]
-        rng_key, sampling_key = jax.random.split(rng_key, 2)
+    def sample_fn(rng_key, size, dtype, mean, cov):
         sample = jax.random.multivariate_normal(
-            sampling_key, mean, cov, shape=size, dtype=dtype, method=op.method
+            rng_key, mean, cov, shape=size, dtype=dtype, method=op.method
         )
-        rng["jax_state"] = rng_key
-        return (rng, sample)
+        return sample
 
     return sample_fn
 
@@ -191,12 +193,9 @@ def jax_sample_fn_bernoulli(op, node):
     """JAX implementation of `BernoulliRV`."""
 
     # We need a separate dispatch, because there is no dtype argument for Bernoulli in JAX
-    def sample_fn(rng, size, dtype, p):
-        rng_key = rng["jax_state"]
-        rng_key, sampling_key = jax.random.split(rng_key, 2)
-        sample = jax.random.bernoulli(sampling_key, p, shape=size)
-        rng["jax_state"] = rng_key
-        return (rng, sample)
+    def sample_fn(rng_key, size, dtype, p):
+        sample = jax.random.bernoulli(rng_key, p, shape=size)
+        return sample
 
     return sample_fn
 
@@ -206,14 +205,10 @@ def jax_sample_fn_categorical(op, node):
     """JAX implementation of `CategoricalRV`."""
 
     # We need a separate dispatch because Categorical expects logits in JAX
-    def sample_fn(rng, size, dtype, p):
-        rng_key = rng["jax_state"]
-        rng_key, sampling_key = jax.random.split(rng_key, 2)
-
+    def sample_fn(rng_key, size, dtype, p):
         logits = jax.scipy.special.logit(p)
-        sample = jax.random.categorical(sampling_key, logits=logits, shape=size)
-        rng["jax_state"] = rng_key
-        return (rng, sample)
+        sample = jax.random.categorical(rng_key, logits=logits, shape=size)
+        return sample
 
     return sample_fn
 
@@ -233,15 +228,10 @@ def jax_sample_fn_uniform(op, node):
         name = "randint"
     jax_op = getattr(jax.random, name)
 
-    def sample_fn(rng, size, dtype, *parameters):
-        rng_key = rng["jax_state"]
-        rng_key, sampling_key = jax.random.split(rng_key, 2)
+    def sample_fn(rng_key, size, dtype, *parameters):
         minval, maxval = parameters
-        sample = jax_op(
-            sampling_key, shape=size, dtype=dtype, minval=minval, maxval=maxval
-        )
-        rng["jax_state"] = rng_key
-        return (rng, sample)
+        sample = jax_op(rng_key, shape=size, dtype=dtype, minval=minval, maxval=maxval)
+        return sample
 
     return sample_fn
 
@@ -258,14 +248,11 @@ def jax_sample_fn_shape_scale(op, node):
     name = op.name
     jax_op = getattr(jax.random, name)
 
-    def sample_fn(rng, size, dtype, shape, scale):
-        rng_key = rng["jax_state"]
-        rng_key, sampling_key = jax.random.split(rng_key, 2)
+    def sample_fn(rng_key, size, dtype, shape, scale):
         if size is None:
             size = jax.numpy.broadcast_arrays(shape, scale)[0].shape
-        sample = jax_op(sampling_key, shape, size, dtype) * scale
-        rng["jax_state"] = rng_key
-        return (rng, sample)
+        sample = jax_op(rng_key, shape, size, dtype) * scale
+        return sample
 
     return sample_fn
 
@@ -274,14 +261,11 @@ def jax_sample_fn_shape_scale(op, node):
 def jax_sample_fn_exponential(op, node):
     """JAX implementation of `ExponentialRV`."""
 
-    def sample_fn(rng, size, dtype, scale):
-        rng_key = rng["jax_state"]
-        rng_key, sampling_key = jax.random.split(rng_key, 2)
+    def sample_fn(rng_key, size, dtype, scale):
         if size is None:
             size = jax.numpy.asarray(scale).shape
-        sample = jax.random.exponential(sampling_key, size, dtype) * scale
-        rng["jax_state"] = rng_key
-        return (rng, sample)
+        sample = jax.random.exponential(rng_key, size, dtype) * scale
+        return sample
 
     return sample_fn
 
@@ -290,14 +274,11 @@ def jax_sample_fn_exponential(op, node):
 def jax_sample_fn_t(op, node):
     """JAX implementation of `StudentTRV`."""
 
-    def sample_fn(rng, size, dtype, df, loc, scale):
-        rng_key = rng["jax_state"]
-        rng_key, sampling_key = jax.random.split(rng_key, 2)
+    def sample_fn(rng_key, size, dtype, df, loc, scale):
         if size is None:
             size = jax.numpy.broadcast_arrays(df, loc, scale)[0].shape
-        sample = loc + jax.random.t(sampling_key, df, size, dtype) * scale
-        rng["jax_state"] = rng_key
-        return (rng, sample)
+        sample = loc + jax.random.t(rng_key, df, size, dtype) * scale
+        return sample
 
     return sample_fn
 
@@ -315,10 +296,7 @@ def jax_funcify_choice(op: ptr.ChoiceWithoutReplacement, node):
             "A default JAX rewrite should have materialized the implicit arange"
         )
 
-    def sample_fn(rng, size, dtype, *parameters):
-        rng_key = rng["jax_state"]
-        rng_key, sampling_key = jax.random.split(rng_key, 2)
-
+    def sample_fn(rng_key, size, dtype, *parameters):
         if op.has_p_param:
             a, p, core_shape = parameters
         else:
@@ -327,9 +305,7 @@ def jax_funcify_choice(op: ptr.ChoiceWithoutReplacement, node):
         core_shape = tuple(np.asarray(core_shape)[(0,) * batch_ndim])
 
         if batch_ndim == 0:
-            sample = jax.random.choice(
-                sampling_key, a, shape=core_shape, replace=False, p=p
-            )
+            sample = jax.random.choice(rng_key, a, shape=core_shape, replace=False, p=p)
 
         else:
             if size is None:
@@ -345,7 +321,7 @@ def jax_funcify_choice(op: ptr.ChoiceWithoutReplacement, node):
             if p is not None:
                 p = jax.numpy.broadcast_to(p, size + p.shape[batch_ndim:])
 
-            batch_sampling_keys = jax.random.split(sampling_key, np.prod(size))
+            batch_sampling_keys = jax.random.split(rng_key, np.prod(size))
 
             # Ravel the batch dimensions because vmap only works along a single axis
             raveled_batch_a = a.reshape((-1,) + a.shape[batch_ndim:])
@@ -366,8 +342,7 @@ def jax_funcify_choice(op: ptr.ChoiceWithoutReplacement, node):
             # Reshape the batch dimensions
             sample = raveled_sample.reshape(size + raveled_sample.shape[1:])
 
-        rng["jax_state"] = rng_key
-        return (rng, sample)
+        return sample
 
     return sample_fn
 
@@ -378,9 +353,7 @@ def jax_sample_fn_permutation(op, node):
 
     batch_ndim = op.batch_ndim(node)
 
-    def sample_fn(rng, size, dtype, *parameters):
-        rng_key = rng["jax_state"]
-        rng_key, sampling_key = jax.random.split(rng_key, 2)
+    def sample_fn(rng_key, size, dtype, *parameters):
         (x,) = parameters
         if batch_ndim:
             # jax.random.permutation has no concept of batch dims
@@ -389,17 +362,16 @@ def jax_sample_fn_permutation(op, node):
             else:
                 x = jax.numpy.broadcast_to(x, size + x.shape[batch_ndim:])
 
-            batch_sampling_keys = jax.random.split(sampling_key, np.prod(size))
+            batch_sampling_keys = jax.random.split(rng_key, np.prod(size))
             raveled_batch_x = x.reshape((-1,) + x.shape[batch_ndim:])
             raveled_sample = jax.vmap(lambda key, x: jax.random.permutation(key, x))(
                 batch_sampling_keys, raveled_batch_x
             )
             sample = raveled_sample.reshape(size + raveled_sample.shape[1:])
         else:
-            sample = jax.random.permutation(sampling_key, x)
+            sample = jax.random.permutation(rng_key, x)
 
-        rng["jax_state"] = rng_key
-        return (rng, sample)
+        return sample
 
     return sample_fn
 
@@ -414,38 +386,45 @@ def jax_sample_fn_binomial(op, node):
 
     from numpyro.distributions.util import binomial
 
-    def sample_fn(rng, size, dtype, n, p):
-        rng_key = rng["jax_state"]
-        rng_key, sampling_key = jax.random.split(rng_key, 2)
-
-        sample = binomial(key=sampling_key, n=n, p=p, shape=size)
-
-        rng["jax_state"] = rng_key
-
-        return (rng, sample)
+    def sample_fn(rng_key, size, dtype, n, p):
+        sample = binomial(key=rng_key, n=n, p=p, shape=size)
+        return sample
 
     return sample_fn
 
 
 @jax_sample_fn.register(ptr.MultinomialRV)
 def jax_sample_fn_multinomial(op, node):
-    if not numpyro_available:
-        raise NotImplementedError(
-            f"No JAX implementation for the given distribution: {op.name}. "
-            "Implementation is available if NumPyro is installed."
-        )
-
-    from numpyro.distributions.util import multinomial
-
-    def sample_fn(rng, size, dtype, n, p):
-        rng_key = rng["jax_state"]
-        rng_key, sampling_key = jax.random.split(rng_key, 2)
-
-        sample = multinomial(key=sampling_key, n=n, p=p, shape=size)
-
-        rng["jax_state"] = rng_key
+    def sample_fn(rng_key, size, dtype, n, p):
+        if size is not None:
+            n = jnp.broadcast_to(n, size)
+            p = jnp.broadcast_to(p, size + jnp.shape(p)[-1:])
 
-        return (rng, sample)
+        else:
+            broadcast_shape = jax.lax.broadcast_shapes(jnp.shape(n), jnp.shape(p)[:-1])
+            n = jnp.broadcast_to(n, broadcast_shape)
+            p = jnp.broadcast_to(p, broadcast_shape + jnp.shape(p)[-1:])
+
+        binom_p = jnp.moveaxis(p, -1, 0)[:-1, ...]
+        sampling_rng = jax.random.split(rng_key, binom_p.shape[0])
+
+        def _binomial_sample_fn(carry, p_rng):
+            s, rho = carry
+            p, rng = p_rng
+            samples = jax.random.binomial(rng, s, p / rho)
+            s = s - samples
+            rho = rho - p
+            return ((s, rho), samples)
+
+        (remain, _), samples = jax.lax.scan(
+            _binomial_sample_fn,
+            (n.astype(np.float64), jnp.ones(binom_p.shape[1:])),
+            (binom_p, sampling_rng),
+        )
+        sample = jnp.concatenate(
+            [jnp.moveaxis(samples, 0, -1), jnp.expand_dims(remain, -1)], axis=-1
+        )
+        return sample
 
     return sample_fn
 
@@ -460,17 +439,12 @@ def jax_sample_fn_vonmises(op, node):
 
     from numpyro.distributions.util import von_mises_centered
 
-    def sample_fn(rng, size, dtype, mu, kappa):
-        rng_key = rng["jax_state"]
-        rng_key, sampling_key = jax.random.split(rng_key, 2)
-
+    def sample_fn(rng_key, size, dtype, mu, kappa):
         sample = von_mises_centered(
-            key=sampling_key, concentration=kappa, shape=size, dtype=dtype
+            key=rng_key, concentration=kappa, shape=size, dtype=dtype
         )
         sample = (sample + mu + np.pi) % (2.0 * np.pi) - np.pi
 
-        rng["jax_state"] = rng_key
-
-        return (rng, sample)
+        return sample
 
     return sample_fn

{pytensor-2.28.3 → pytensor-2.30.0}/pytensor/link/jax/dispatch/scan.py

@@ -29,7 +29,7 @@ def jax_funcify_Scan(op: Scan, **kwargs):
         # Extract JAX scan inputs
         outer_inputs = list(outer_inputs)
         n_steps = outer_inputs[0]  # JAX `length`
-        seqs = op.outer_seqs(outer_inputs)  # JAX `xs`
+        seqs = [seq[:n_steps] for seq in op.outer_seqs(outer_inputs)]  # JAX `xs`
 
         mit_sot_init = []
         for tap, seq in zip(

{pytensor-2.28.3 → pytensor-2.30.0}/pytensor/link/jax/dispatch/shape.py

@@ -4,7 +4,7 @@ from pytensor.graph import Constant
 from pytensor.graph.basic import Apply
 from pytensor.graph.op import Op
 from pytensor.link.jax.dispatch.basic import jax_funcify
-from pytensor.tensor.shape import Reshape, Shape, Shape_i, SpecifyShape, Unbroadcast
+from pytensor.tensor.shape import Reshape, Shape, Shape_i, SpecifyShape
 from pytensor.tensor.type import TensorType
 
 
@@ -104,11 +104,3 @@ def jax_funcify_SpecifyShape(op, node, **kwargs):
         return x
 
     return specifyshape
-
-
-@jax_funcify.register(Unbroadcast)
-def jax_funcify_Unbroadcast(op, **kwargs):
-    def unbroadcast(x):
-        return x
-
-    return unbroadcast

pytensor-2.30.0/pytensor/link/jax/dispatch/signal/__init__.py

@@ -0,0 +1 @@
+import pytensor.link.jax.dispatch.signal.conv

pytensor-2.30.0/pytensor/link/jax/dispatch/signal/conv.py

@@ -0,0 +1,14 @@
+import jax
+
+from pytensor.link.jax.dispatch import jax_funcify
+from pytensor.tensor.signal.conv import Conv1d
+
+
+@jax_funcify.register(Conv1d)
+def jax_funcify_Conv1d(op, node, **kwargs):
+    mode = op.mode
+
+    def conv1d(data, kernel):
+        return jax.numpy.convolve(data, kernel, mode=mode)
+
+    return conv1d

{pytensor-2.28.3 → pytensor-2.30.0}/pytensor/link/jax/dispatch/slinalg.py

@@ -1,3 +1,5 @@
+import warnings
+
 import jax
 
 from pytensor.link.jax.dispatch.basic import jax_funcify
@@ -39,13 +41,29 @@ def jax_funcify_Cholesky(op, **kwargs):
 
 @jax_funcify.register(Solve)
 def jax_funcify_Solve(op, **kwargs):
-    if op.assume_a != "gen" and op.lower:
-        lower = True
+    assume_a = op.assume_a
+    lower = op.lower
+
+    if assume_a == "tridiagonal":
+        # jax.scipy.solve does not yet support tridiagonal matrices
+        # But there's a jax.lax.linalg.tridiaonal_solve we can use instead.
+        def solve(a, b):
+            dl = jax.numpy.diagonal(a, offset=-1, axis1=-2, axis2=-1)
+            d = jax.numpy.diagonal(a, offset=0, axis1=-2, axis2=-1)
+            du = jax.numpy.diagonal(a, offset=1, axis1=-2, axis2=-1)
+            return jax.lax.linalg.tridiagonal_solve(dl, d, du, b, lower=lower)
+
     else:
-        lower = False
+        if assume_a not in ("gen", "sym", "her", "pos"):
+            warnings.warn(
+                f"JAX solve does not support assume_a={op.assume_a}. Defaulting to assume_a='gen'.\n"
+                f"If appropriate, you may want to set assume_a to one of 'sym', 'pos', 'her' or 'tridiagonal' to improve performance.",
+                UserWarning,
+            )
+            assume_a = "gen"
 
-    def solve(a, b, lower=lower):
-        return jax.scipy.linalg.solve(a, b, lower=lower)
+        def solve(a, b):
+            return jax.scipy.linalg.solve(a, b, lower=lower, assume_a=assume_a)
 
     return solve
 

{pytensor-2.28.3 → pytensor-2.30.0}/pytensor/link/numba/dispatch/__init__.py

@@ -9,9 +9,11 @@ import pytensor.link.numba.dispatch.nlinalg
 import pytensor.link.numba.dispatch.random
 import pytensor.link.numba.dispatch.scan
 import pytensor.link.numba.dispatch.scalar
+import pytensor.link.numba.dispatch.signal
 import pytensor.link.numba.dispatch.slinalg
 import pytensor.link.numba.dispatch.sparse
 import pytensor.link.numba.dispatch.subtensor
 import pytensor.link.numba.dispatch.tensor_basic
 
+
 # isort: on