Repository: pymc-devs/pytensor

Issue:
Implement gradient for QR decomposition
### Description

QR is one of the few remaining linalg ops that is missing a gradient. JAX code for the jvp is [here](https://github.com/jax-ml/jax/blob/73fa0f48cb0081fc69cb82dffe6ceb5433cdc446/jax/_src/lax/linalg.py#L1954), whic also includes [this derivation](https://j-towns.github.io/papers/qr-derivative.pdf). [This paper ](https://www.tandfonline.com/doi/abs/10.1080/10556788.2011.610454) also claims to derive the gradients for QR, but I find it unreadable.

Relatedly but perhaps worthy of a separate issue, [this paper ](https://arxiv.org/abs/1710.08717) derives gradients for the LQ decomposition, $$A = LQ$$, where $L$ is lower triangular and $Q$ is orthonormal ($$Q^TQ=I$$.)  Compare this to QR, which gives you $$A = QR$$, where $$Q$$ is again orthonormal, but $$R$$ is upper triangular, and you see why I mention it in this issue. It wouldn't be hard to offer LQ as well.

Discussion:
@educhesne: I went for the backward-mode version of the solution implemented in jax. So the gradient is only valid when R is square and full-rank.
A note about the implementation: at some point the inverse of R is needed, and I used matrix_inverse rather than solve_triangular because I assumed that no dependency on scipy was allowed in nlinalg; was I correct ? 
(first PR to pytensor here; any comment appreciated)

> JAX code for the jvp is [here](https://github.com/jax-ml/jax/blob/73fa0f48cb0081fc69cb82dffe6ceb5433cdc446/jax/_src/lax/linalg.py#L1954), whic also includes [this derivation](https://j-towns.github.io/papers/qr-derivative.pdf).

This solution requires R to be square and full-rank for the gradient to be valid.


> [This paper ](https://www.tandfonline.com/doi/abs/10.1080/10556788.2011.610454) also claims to derive the gradients for QR, but I find it unreadable.

As far as I can tell, this solution only works on square input (based on some verify_grad I ran; I can't really pinpoint which part of the derivation assumes it...)


> Relatedly but perhaps worthy of a separate issue, [this paper ](https://arxiv.org/abs/1710.08717) derives gradients for the LQ decomposition

I couldn't really make sense of this derivation... maybe there are some typos in it ?


> It wouldn't be hard to offer LQ as well.

I didn't do it, although if you want I can add LQ as a QR decomposition with a couple of transpose

@qiyang-ustc: I the past I always refer to this blog. will it help [https://giggleliu.github.io/posts/2019-04-02-einsumbp/](https://giggleliu.github.io/posts/2019-04-02-einsumbp/)?
@jessegrabowski: Very nice reference! Maybe the einsum rule could improve on what we're doing now (which I think just autodiffing through whatever graph einsum comes up with)@educhesne: The derivation is much clearer in that blog ! Thanks ! Besides in the paper they reference, there is an expression for QR gradient when m<n (where QR is a m x n matrix); that was not covered in the previous references. I'm going to look into it.@educhesne: I updated the PR with the gradient expression covering the others shapes from that [paper](https://arxiv.org/abs/1903.09650) (I think there is actually a tiny mistake in it, but I reckon it is corrected in the PR)

Solution patch:
diff --git a/pytensor/graph/op.py b/pytensor/graph/op.py
index 690bb44df..b2d70d582 100644
--- a/pytensor/graph/op.py
+++ b/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,
diff --git a/pytensor/link/jax/dispatch/random.py b/pytensor/link/jax/dispatch/random.py
index 8a33dfac1..678b2dc48 100644
--- a/pytensor/link/jax/dispatch/random.py
+++ b/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
diff --git a/pytensor/link/numba/dispatch/basic.py b/pytensor/link/numba/dispatch/basic.py
index 19e91e5f8..0958e5e77 100644
--- a/pytensor/link/numba/dispatch/basic.py
+++ b/pytensor/link/numba/dispatch/basic.py
@@ -33,6 +33,7 @@ from pytensor.link.utils import (
 from pytensor.scalar.basic import ScalarType
 from pytensor.scalar.math import Softplus
 from pytensor.sparse import SparseTensorType
+from pytensor.tensor.basic import Nonzero
 from pytensor.tensor.blas import BatchedDot
 from pytensor.tensor.math import Dot
 from pytensor.tensor.shape import Reshape, Shape, Shape_i, SpecifyShape
@@ -657,3 +658,15 @@ def numba_funcify_IfElse(op, **kwargs):
             return res[0]
 
     return ifelse
+
+
+@numba_funcify.register(Nonzero)
+def numba_funcify_Nonzero(op, node, **kwargs):
+    @numba_njit
+    def nonzero(a):
+        result_tuple = np.nonzero(a)
+        if a.ndim == 1:
+            return result_tuple[0]
+        return list(result_tuple)
+
+    return nonzero
diff --git a/pytensor/link/numba/dispatch/slinalg.py b/pytensor/link/numba/dispatch/slinalg.py
index 700bd57d4..f02512ca5 100644
--- a/pytensor/link/numba/dispatch/slinalg.py
+++ b/pytensor/link/numba/dispatch/slinalg.py
@@ -24,6 +24,19 @@ from pytensor.tensor.slinalg import (
     Solve,
     SolveTriangular,
 )
+from pytensor.tensor.type import complex_dtypes
+
+
+_COMPLEX_DTYPE_NOT_SUPPORTED_MSG = (
+    "Complex dtype for {op} not supported in numba mode. "
+    "If you need this functionality, please open an issue at: https://github.com/pymc-devs/pytensor"
+)
+
+
+@numba_basic.numba_njit(inline="always")
+def _copy_to_fortran_order_even_if_1d(x):
+    # Numba's _copy_to_fortran_order doesn't do anything for vectors
+    return x.copy() if x.ndim == 1 else _copy_to_fortran_order(x)
 
 
 @numba_basic.numba_njit(inline="always")
@@ -120,6 +133,9 @@ def solve_triangular_impl(A, B, trans, lower, unit_diagonal, b_ndim, overwrite_b
     dtype = A.dtype
     w_type = _get_underlying_float(dtype)
     numba_trtrs = _LAPACK().numba_xtrtrs(dtype)
+    if isinstance(dtype, types.Complex):
+        # If you want to make this work with complex numbers make sure you handle the c_contiguous trick correctly
+        raise TypeError("This function is not expected to work with complex numbers")
 
     def impl(A, B, trans, lower, unit_diagonal, b_ndim, overwrite_b):
         _N = np.int32(A.shape[-1])
@@ -129,21 +145,23 @@ def solve_triangular_impl(A, B, trans, lower, unit_diagonal, b_ndim, overwrite_b
         # could potentially be 3d (it didn't understand b_ndim was always equal to B.ndim)
         B_is_1d = B.ndim == 1
 
-        # This will only copy if A is not already fortran contiguous
-        A_f = np.asfortranarray(A)
+        if A.flags.f_contiguous or (A.flags.c_contiguous and trans in (0, 1)):
+            A_f = A
+            if A.flags.c_contiguous:
+                # An upper/lower triangular c_contiguous is the same as a lower/upper triangular f_contiguous
+                # Is this valid for complex matrices that were .conj().mT by PyTensor?
+                lower = not lower
+                trans = 1 - trans
+        else:
+            A_f = np.asfortranarray(A)
 
-        if overwrite_b:
-            if B_is_1d:
-                B_copy = np.expand_dims(B, -1)
-            else:
-                # This *will* allow inplace destruction of B, but only if it is already fortran contiguous.
-                # Otherwise, there's no way to get around the need to copy the data before going into TRTRS
-                B_copy = np.asfortranarray(B)
+        if overwrite_b and B.flags.f_contiguous:
+            B_copy = B
         else:
-            if B_is_1d:
-                B_copy = np.copy(np.expand_dims(B, -1))
-            else:
-                B_copy = _copy_to_fortran_order(B)
+            B_copy = _copy_to_fortran_order_even_if_1d(B)
+
+        if B_is_1d:
+            B_copy = np.expand_dims(B_copy, -1)
 
         NRHS = 1 if B_is_1d else int(B_copy.shape[-1])
 
@@ -188,9 +206,9 @@ def numba_funcify_SolveTriangular(op, node, **kwargs):
     b_ndim = op.b_ndim
 
     dtype = node.inputs[0].dtype
-    if str(dtype).startswith("complex"):
+    if dtype in complex_dtypes:
         raise NotImplementedError(
-            "Complex inputs not currently supported by solve_triangular in Numba mode"
+            _COMPLEX_DTYPE_NOT_SUPPORTED_MSG.format(op="Solve Triangular")
         )
 
     @numba_basic.numba_njit(inline="always")
@@ -247,10 +265,10 @@ def cholesky_impl(A, lower=0, overwrite_a=False, check_finite=True):
         LDA = val_to_int_ptr(_N)
         INFO = val_to_int_ptr(0)
 
-        if not overwrite_a:
-            A_copy = _copy_to_fortran_order(A)
-        else:
+        if overwrite_a and A.flags.f_contiguous:
             A_copy = A
+        else:
+            A_copy = _copy_to_fortran_order(A)
 
         numba_potrf(
             UPLO,
@@ -283,15 +301,13 @@ def numba_funcify_Cholesky(op, node, **kwargs):
     In particular, the `inplace` argument is not supported, which is why we choose to implement our own version.
     """
     lower = op.lower
-    overwrite_a = False
+    overwrite_a = op.overwrite_a
     check_finite = op.check_finite
     on_error = op.on_error
 
     dtype = node.inputs[0].dtype
-    if str(dtype).startswith("complex"):
-        raise NotImplementedError(
-            "Complex inputs not currently supported by cholesky in Numba mode"
-        )
+    if dtype in complex_dtypes:
+        raise NotImplementedError(_COMPLEX_DTYPE_NOT_SUPPORTED_MSG.format(op=op))
 
     @numba_basic.numba_njit(inline="always")
     def nb_cholesky(a):
@@ -497,10 +513,10 @@ def getrf_impl(
     ) -> tuple[np.ndarray, np.ndarray, int]:
         _M, _N = np.int32(A.shape[-2:])  # type: ignore
 
-        if not overwrite_a:
-            A_copy = _copy_to_fortran_order(A)
-        else:
+        if overwrite_a and A.flags.f_contiguous:
             A_copy = A
+        else:
+            A_copy = _copy_to_fortran_order(A)
 
         M = val_to_int_ptr(_M)  # type: ignore
         N = val_to_int_ptr(_N)  # type: ignore
@@ -545,10 +561,10 @@ def getrs_impl(
 
         B_is_1d = B.ndim == 1
 
-        if not overwrite_b:
-            B_copy = _copy_to_fortran_order(B)
-        else:
+        if overwrite_b and B.flags.f_contiguous:
             B_copy = B
+        else:
+            B_copy = _copy_to_fortran_order_even_if_1d(B)
 
         if B_is_1d:
             B_copy = np.expand_dims(B_copy, -1)
@@ -576,7 +592,7 @@ def getrs_impl(
         )
 
         if B_is_1d:
-            return B_copy[..., 0], int_ptr_to_val(INFO)
+            B_copy = B_copy[..., 0]
 
         return B_copy, int_ptr_to_val(INFO)
 
@@ -632,6 +648,11 @@ def solve_gen_impl(
         _N = np.int32(A.shape[-1])
         _solve_check_input_shapes(A, B)
 
+        if overwrite_a and A.flags.c_contiguous:
+            # Work with the transposed system to avoid copying A
+            A = A.T
+            transposed = not transposed
+
         order = "I" if transposed else "1"
         norm = _xlange(A, order=order)
 
@@ -681,19 +702,23 @@ def sysv_impl(
         _LDA, _N = np.int32(A.shape[-2:])  # type: ignore
         _solve_check_input_shapes(A, B)
 
-        if not overwrite_a:
-            A_copy = _copy_to_fortran_order(A)
-        else:
+        if overwrite_a and (A.flags.f_contiguous or A.flags.c_contiguous):
             A_copy = A
+            if A.flags.c_contiguous:
+                # An upper/lower triangular c_contiguous is the same as a lower/upper triangular f_contiguous
+                lower = not lower
+        else:
+            A_copy = _copy_to_fortran_order(A)
 
         B_is_1d = B.ndim == 1
 
-        if not overwrite_b:
-            B_copy = _copy_to_fortran_order(B)
-        else:
+        if overwrite_b and B.flags.f_contiguous:
             B_copy = B
+        else:
+            B_copy = _copy_to_fortran_order_even_if_1d(B)
+
         if B_is_1d:
-            B_copy = np.asfortranarray(np.expand_dims(B_copy, -1))
+            B_copy = np.expand_dims(B_copy, -1)
 
         NRHS = 1 if B_is_1d else int(B.shape[-1])
 
@@ -864,7 +889,7 @@ def _posv(
     overwrite_b: bool,
     check_finite: bool,
     transposed: bool,
-) -> tuple[np.ndarray, int]:
+) -> tuple[np.ndarray, np.ndarray, int]:
     """
     Placeholder for solving a linear system with a positive-definite matrix; used by linalg.solve.
     """
@@ -881,7 +906,8 @@ def posv_impl(
     check_finite: bool,
     transposed: bool,
 ) -> Callable[
-    [np.ndarray, np.ndarray, bool, bool, bool, bool, bool], tuple[np.ndarray, int]
+    [np.ndarray, np.ndarray, bool, bool, bool, bool, bool],
+    tuple[np.ndarray, np.ndarray, int],
 ]:
     ensure_lapack()
     _check_scipy_linalg_matrix(A, "solve")
@@ -898,22 +924,25 @@ def posv_impl(
         overwrite_b: bool,
         check_finite: bool,
         transposed: bool,
-    ) -> tuple[np.ndarray, int]:
+    ) -> tuple[np.ndarray, np.ndarray, int]:
         _solve_check_input_shapes(A, B)
 
         _N = np.int32(A.shape[-1])
 
-        if not overwrite_a:
-            A_copy = _copy_to_fortran_order(A)
-        else:
+        if overwrite_a and (A.flags.f_contiguous or A.flags.c_contiguous):
             A_copy = A
+            if A.flags.c_contiguous:
+                # An upper/lower triangular c_contiguous is the same as a lower/upper triangular f_contiguous
+                lower = not lower
+        else:
+            A_copy = _copy_to_fortran_order(A)
 
         B_is_1d = B.ndim == 1
 
-        if not overwrite_b:
-            B_copy = _copy_to_fortran_order(B)
-        else:
+        if overwrite_b and B.flags.f_contiguous:
             B_copy = B
+        else:
+            B_copy = _copy_to_fortran_order_even_if_1d(B)
 
         if B_is_1d:
             B_copy = np.expand_dims(B_copy, -1)
@@ -939,8 +968,9 @@ def posv_impl(
         )
 
         if B_is_1d:
-            return B_copy[..., 0], int_ptr_to_val(INFO)
-        return B_copy, int_ptr_to_val(INFO)
+            B_copy = B_copy[..., 0]
+
+        return A_copy, B_copy, int_ptr_to_val(INFO)
 
     return impl
 
@@ -1041,10 +1071,12 @@ def solve_psd_impl(
     ) -> np.ndarray:
         _solve_check_input_shapes(A, B)
 
-        x, info = _posv(A, B, lower, overwrite_a, overwrite_b, check_finite, transposed)
+        C, x, info = _posv(
+            A, B, lower, overwrite_a, overwrite_b, check_finite, transposed
+        )
         _solve_check(A.shape[-1], info)
 
-        rcond, info = _pocon(x, _xlange(A))
+        rcond, info = _pocon(C, _xlange(A))
         _solve_check(A.shape[-1], info=info, lamch=True, rcond=rcond)
 
         return x
@@ -1062,10 +1094,8 @@ def numba_funcify_Solve(op, node, **kwargs):
     transposed = False  # TODO: Solve doesnt currently allow the transposed argument
 
     dtype = node.inputs[0].dtype
-    if str(dtype).startswith("complex"):
-        raise NotImplementedError(
-            "Complex inputs not currently supported by solve in Numba mode"
-        )
+    if dtype in complex_dtypes:
+        raise NotImplementedError(_COMPLEX_DTYPE_NOT_SUPPORTED_MSG.format(op=op))
 
     if assume_a == "gen":
         solve_fn = _solve_gen
@@ -1102,12 +1132,15 @@ def numba_funcify_Solve(op, node, **kwargs):
     return solve
 
 
-def _cho_solve(A_and_lower, B, overwrite_a=False, overwrite_b=False, check_finite=True):
+def _cho_solve(
+    C: np.ndarray, B: np.ndarray, lower: bool, overwrite_b: bool, check_finite: bool
+):
     """
     Solve a positive-definite linear system using the Cholesky decomposition.
     """
-    A, lower = A_and_lower
-    return linalg.cho_solve((A, lower), B)
+    return linalg.cho_solve(
+        (C, lower), b=B, overwrite_b=overwrite_b, check_finite=check_finite
+    )
 
 
 @overload(_cho_solve)
@@ -1123,13 +1156,22 @@ def cho_solve_impl(C, B, lower=False, overwrite_b=False, check_finite=True):
         _solve_check_input_shapes(C, B)
 
         _N = np.int32(C.shape[-1])
-        C_copy = _copy_to_fortran_order(C)
+        if C.flags.f_contiguous or C.flags.c_contiguous:
+            C_f = C
+            if C.flags.c_contiguous:
+                # An upper/lower triangular c_contiguous is the same as a lower/upper triangular f_contiguous
+                lower = not lower
+        else:
+            C_f = np.asfortranarray(C)
+
+        if overwrite_b and B.flags.f_contiguous:
+            B_copy = B
+        else:
+            B_copy = _copy_to_fortran_order_even_if_1d(B)
 
         B_is_1d = B.ndim == 1
         if B_is_1d:
-            B_copy = np.asfortranarray(np.expand_dims(B, -1))
-        else:
-            B_copy = _copy_to_fortran_order(B)
+            B_copy = np.expand_dims(B_copy, -1)
 
         NRHS = 1 if B_is_1d else int(B.shape[-1])
 
@@ -1144,16 +1186,18 @@ def cho_solve_impl(C, B, lower=False, overwrite_b=False, check_finite=True):
             UPLO,
             N,
             NRHS,
-            C_copy.view(w_type).ctypes,
+            C_f.view(w_type).ctypes,
             LDA,
             B_copy.view(w_type).ctypes,
             LDB,
             INFO,
         )
 
+        _solve_check(_N, int_ptr_to_val(INFO))
+
         if B_is_1d:
-            return B_copy[..., 0], int_ptr_to_val(INFO)
-        return B_copy, int_ptr_to_val(INFO)
+            return B_copy[..., 0]
+        return B_copy
 
     return impl
 
@@ -1165,10 +1209,8 @@ def numba_funcify_CholeskySolve(op, node, **kwargs):
     check_finite = op.check_finite
 
     dtype = node.inputs[0].dtype
-    if str(dtype).startswith("complex"):
-        raise NotImplementedError(
-            "Complex inputs not currently supported by cho_solve in Numba mode"
-        )
+    if dtype in complex_dtypes:
+        raise NotImplementedError(_COMPLEX_DTYPE_NOT_SUPPORTED_MSG.format(op=op))
 
     @numba_basic.numba_njit(inline="always")
     def cho_solve(c, b):
@@ -1182,16 +1224,8 @@ def numba_funcify_CholeskySolve(op, node, **kwargs):
                     "Non-numeric values (nan or inf) in input b to cho_solve"
                 )
 
-        res, info = _cho_solve(
+        return _cho_solve(
             c, b, lower=lower, overwrite_b=overwrite_b, check_finite=check_finite
         )
 
-        if info < 0:
-            raise np.linalg.LinAlgError("Illegal values found in input to cho_solve")
-        elif info > 0:
-            raise np.linalg.LinAlgError(
-                "Matrix is not positive definite in input to cho_solve"
-            )
-        return res
-
     return cho_solve
diff --git a/pytensor/tensor/basic.py b/pytensor/tensor/basic.py
index 4a4ae5f15..6bcb084f4 100644
--- a/pytensor/tensor/basic.py
+++ b/pytensor/tensor/basic.py
@@ -3215,13 +3215,29 @@ class ARange(Op):
         self.dtype = dtype
 
     def make_node(self, start, stop, step):
+        from math import ceil
+
         start, stop, step = map(as_tensor_variable, (start, stop, step))
+
         assert start.ndim == 0
         assert stop.ndim == 0
         assert step.ndim == 0
 
+        # if it is possible to directly determine the shape i.e static shape is present, we find it.
+        if (
+            isinstance(start, TensorConstant)
+            and isinstance(stop, TensorConstant)
+            and isinstance(step, TensorConstant)
+        ):
+            length = max(
+                ceil((float(stop.data) - float(start.data)) / float(step.data)), 0
+            )
+            shape = (length,)
+        else:
+            shape = (None,)
+
         inputs = [start, stop, step]
-        outputs = [tensor(dtype=self.dtype, shape=(None,))]
+        outputs = [tensor(dtype=self.dtype, shape=shape)]
 
         return Apply(self, inputs, outputs)
 
diff --git a/pytensor/tensor/math.py b/pytensor/tensor/math.py
index 2ba2ab158..714f597b3 100644
--- a/pytensor/tensor/math.py
+++ b/pytensor/tensor/math.py
@@ -2453,6 +2453,11 @@ def kv(v, x):
     return kve(v, x) * exp(-x)
 
 
+def kn(n, x):
+    """Modified Bessel function of the second kind of integer order v."""
+    return kv(n, x)
+
+
 @scalar_elemwise
 def sigmoid(x):
     """Logistic sigmoid function (1 / (1 + exp(-x)), also known as expit or inverse logit"""
@@ -4337,6 +4342,7 @@ __all__ = [
     "i1",
     "iv",
     "ive",
+    "kn",
     "kv",
     "kve",
     "sigmoid",
diff --git a/pytensor/tensor/nlinalg.py b/pytensor/tensor/nlinalg.py
index ee33f6533..8fff2a2f5 100644
--- a/pytensor/tensor/nlinalg.py
+++ b/pytensor/tensor/nlinalg.py
@@ -5,12 +5,15 @@ from typing import Literal, cast
 
 import numpy as np
 
+import pytensor.tensor as pt
 from pytensor import scalar as ps
 from pytensor.compile.builders import OpFromGraph
 from pytensor.gradient import DisconnectedType
 from pytensor.graph.basic import Apply
 from pytensor.graph.op import Op
+from pytensor.ifelse import ifelse
 from pytensor.npy_2_compat import normalize_axis_tuple
+from pytensor.raise_op import Assert
 from pytensor.tensor import TensorLike
 from pytensor.tensor import basic as ptb
 from pytensor.tensor import math as ptm
@@ -512,6 +515,80 @@ class QRFull(Op):
         else:
             outputs[0][0] = res
 
+    def L_op(self, inputs, outputs, output_grads):
+        """
+        Reverse-mode gradient of the QR function.
+
+        References
+        ----------
+        .. [1] Jinguo Liu. "Linear Algebra Autodiff (complex valued)", blog post https://giggleliu.github.io/posts/2019-04-02-einsumbp/
+        .. [2] Hai-Jun Liao, Jin-Guo Liu, Lei Wang, Tao Xiang. "Differentiable Programming Tensor Networks", arXiv:1903.09650v2
+        """
+
+        from pytensor.tensor.slinalg import solve_triangular
+
+        (A,) = (cast(ptb.TensorVariable, x) for x in inputs)
+        m, n = A.shape
+
+        def _H(x: ptb.TensorVariable):
+            return x.conj().mT
+
+        def _copyltu(x: ptb.TensorVariable):
+            return ptb.tril(x, k=0) + _H(ptb.tril(x, k=-1))
+
+        if self.mode == "raw":
+            raise NotImplementedError("Gradient of qr not implemented for mode=raw")
+
+        elif self.mode == "r":
+            # We need all the components of the QR to compute the gradient of A even if we only
+            # use the upper triangular component in the cost function.
+            Q, R = qr(A, mode="reduced")
+            dQ = Q.zeros_like()
+            dR = cast(ptb.TensorVariable, output_grads[0])
+
+        else:
+            Q, R = (cast(ptb.TensorVariable, x) for x in outputs)
+            if self.mode == "complete":
+                qr_assert_op = Assert(
+                    "Gradient of qr not implemented for m x n matrices with m > n and mode=complete"
+                )
+                R = qr_assert_op(R, ptm.le(m, n))
+
+            new_output_grads = []
+            is_disconnected = [
+                isinstance(x.type, DisconnectedType) for x in output_grads
+            ]
+            if all(is_disconnected):
+                # This should never be reached by Pytensor
+                return [DisconnectedType()()]  # pragma: no cover
+
+            for disconnected, output_grad, output in zip(
+                is_disconnected, output_grads, [Q, R], strict=True
+            ):
+                if disconnected:
+                    new_output_grads.append(output.zeros_like())
+                else:
+                    new_output_grads.append(output_grad)
+
+            (dQ, dR) = (cast(ptb.TensorVariable, x) for x in new_output_grads)
+
+        # gradient expression when m >= n
+        M = R @ _H(dR) - _H(dQ) @ Q
+        K = dQ + Q @ _copyltu(M)
+        A_bar_m_ge_n = _H(solve_triangular(R, _H(K)))
+
+        # gradient expression when m < n
+        Y = A[:, m:]
+        U = R[:, :m]
+        dU, dV = dR[:, :m], dR[:, m:]
+        dQ_Yt_dV = dQ + Y @ _H(dV)
+        M = U @ _H(dU) - _H(dQ_Yt_dV) @ Q
+        X_bar = _H(solve_triangular(U, _H(dQ_Yt_dV + Q @ _copyltu(M))))
+        Y_bar = Q @ dV
+        A_bar_m_lt_n = pt.concatenate([X_bar, Y_bar], axis=1)
+
+        return [ifelse(ptm.ge(m, n), A_bar_m_ge_n, A_bar_m_lt_n)]
+
 
 def qr(a, mode="reduced"):
     """
diff --git a/pytensor/tensor/sharedvar.py b/pytensor/tensor/sharedvar.py
index 7ce15c272..f193cf8dc 100644
--- a/pytensor/tensor/sharedvar.py
+++ b/pytensor/tensor/sharedvar.py
@@ -3,7 +3,6 @@ import warnings
 import numpy as np
 
 from pytensor.compile import SharedVariable, shared_constructor
-from pytensor.tensor import _get_vector_length
 from pytensor.tensor.type import TensorType
 from pytensor.tensor.variable import TensorVariable
 
@@ -51,11 +50,6 @@ class TensorSharedVariable(SharedVariable, TensorVariable):
             self.container.value = 0 * self.container.value
 
 
-@_get_vector_length.register(TensorSharedVariable)
-def _get_vector_length_TensorSharedVariable(var_inst, var):
-    return len(var.get_value(borrow=True))
-
-
 @shared_constructor.register(np.ndarray)
 def tensor_constructor(
     value,


Test patch:
diff --git a/tests/link/jax/test_random.py b/tests/link/jax/test_random.py
index 2a6ebca0a..183b629f7 100644
--- a/tests/link/jax/test_random.py
+++ b/tests/link/jax/test_random.py
@@ -703,14 +703,15 @@ def test_beta_binomial():
     )
 
 
-@pytest.mark.skipif(
-    not numpyro_available, reason="Multinomial dispatch requires numpyro"
-)
 def test_multinomial():
     rng = shared(np.random.default_rng(123))
+
+    # test with 'size' argument and n.shape == p.shape[:-1]
     n = np.array([10, 40])
     p = np.array([[0.3, 0.7, 0.0], [0.1, 0.4, 0.5]])
-    g = pt.random.multinomial(n, p, size=(10_000, 2), rng=rng)
+    size = (10_000, 2)
+
+    g = pt.random.multinomial(n, p, size=size, rng=rng)
     g_fn = compile_random_function([], g, mode="JAX")
     samples = g_fn()
     np.testing.assert_allclose(samples.mean(axis=0), n[..., None] * p, rtol=0.1)
@@ -718,6 +719,20 @@ def test_multinomial():
         samples.std(axis=0), np.sqrt(n[..., None] * p * (1 - p)), rtol=0.1
     )
 
+    # test with no 'size' argument and no static shape
+    n = np.broadcast_to(np.array([10, 40]), size)
+    p = np.array([[0.3, 0.7, 0.0], [0.1, 0.4, 0.5]])
+    pt_n = pt.matrix("n")
+    pt_p = pt.matrix("p")
+
+    g = pt.random.multinomial(pt_n, pt_p, rng=rng, size=None)
+    g_fn = compile_random_function([pt_n, pt_p], g, mode="JAX")
+    samples = g_fn(n, p)
+    np.testing.assert_allclose(samples.mean(axis=0), n[0, :, None] * p, rtol=0.1)
+    np.testing.assert_allclose(
+        samples.std(axis=0), np.sqrt(n[0, :, None] * p * (1 - p)), rtol=0.1
+    )
+
 
 @pytest.mark.skipif(not numpyro_available, reason="VonMises dispatch requires numpyro")
 def test_vonmises_mu_outside_circle():
@@ -796,7 +811,7 @@ def test_random_custom_implementation():
     @jax_sample_fn.register(CustomRV)
     def jax_sample_fn_custom(op, node):
         def sample_fn(rng, size, dtype, *parameters):
-            return (rng, 0)
+            return 0
 
         return sample_fn
 
diff --git a/tests/link/numba/test_basic.py b/tests/link/numba/test_basic.py
index 654cbe7bd..5fb530985 100644
--- a/tests/link/numba/test_basic.py
+++ b/tests/link/numba/test_basic.py
@@ -7,6 +7,7 @@ from unittest import mock
 import numpy as np
 import pytest
 
+from pytensor.compile import SymbolicInput
 from tests.tensor.test_math_scipy import scipy
 
 
@@ -120,6 +121,7 @@ opts = RewriteDatabaseQuery(
 numba_mode = Mode(
     NumbaLinker(), opts.including("numba", "local_useless_unbatched_blockwise")
 )
+numba_inplace_mode = numba_mode.including("inplace")
 py_mode = Mode("py", opts)
 
 rng = np.random.default_rng(42849)
@@ -261,7 +263,11 @@ def compare_numba_and_py(
                 x, y
             )
 
-    if any(inp.owner is not None for inp in graph_inputs):
+    if any(
+        inp.owner is not None
+        for inp in graph_inputs
+        if not isinstance(inp, SymbolicInput)
+    ):
         raise ValueError("Inputs must be root variables")
 
     pytensor_py_fn = function(
@@ -287,7 +293,6 @@ def compare_numba_and_py(
     )
     test_inputs_copy = (inp.copy() for inp in test_inputs) if inplace else test_inputs
     numba_res = pytensor_numba_fn(*test_inputs_copy)
-
     if isinstance(graph_outputs, tuple | list):
         for j, p in zip(numba_res, py_res, strict=True):
             assert_fn(j, p)
@@ -893,3 +898,17 @@ def test_function_overhead(mode, benchmark):
     assert np.sum(fn(test_x)) == 1000
 
     benchmark(fn, test_x)
+
+
+@pytest.mark.parametrize(
+    "input_data",
+    [np.array([1, 0, 3]), np.array([[0, 1], [2, 0]]), np.array([[0, 0], [0, 0]])],
+)
+def test_Nonzero(input_data):
+    a = pt.tensor("a", shape=(None,) * input_data.ndim)
+
+    graph_outputs = pt.nonzero(a)
+
+    compare_numba_and_py(
+        graph_inputs=[a], graph_outputs=graph_outputs, test_inputs=[input_data]
+    )
diff --git a/tests/link/numba/test_slinalg.py b/tests/link/numba/test_slinalg.py
index defbcf6c8..51c367bbc 100644
--- a/tests/link/numba/test_slinalg.py
+++ b/tests/link/numba/test_slinalg.py
@@ -1,17 +1,15 @@
 import re
-from functools import partial
 from typing import Literal
 
 import numpy as np
 import pytest
-from numpy.testing import assert_allclose
+import scipy
 
 import pytensor
 import pytensor.tensor as pt
-from pytensor import config
-from pytensor.tensor.slinalg import SolveTriangular
-from tests import unittest_tools as utt
-from tests.link.numba.test_basic import compare_numba_and_py
+from pytensor import In, config
+from pytensor.tensor.slinalg import Cholesky, CholeskySolve, Solve, SolveTriangular
+from tests.link.numba.test_basic import compare_numba_and_py, numba_inplace_mode
 
 
 numba = pytest.importorskip("numba")
@@ -21,250 +19,6 @@ floatX = config.floatX
 rng = np.random.default_rng(42849)
 
 
-def transpose_func(x, trans):
-    if trans == 0:
-        return x
-    if trans == 1:
-        return x.T
-    if trans == 2:
-        return x.conj().T
-
-
-@pytest.mark.parametrize(
-    "b_shape",
-    [(5, 1), (5, 5), (5,)],
-    ids=["b_col_vec", "b_matrix", "b_vec"],
-)
-@pytest.mark.parametrize("lower", [True, False], ids=["lower=True", "lower=False"])
-@pytest.mark.parametrize("trans", [0, 1, 2], ids=["trans=N", "trans=C", "trans=T"])
-@pytest.mark.parametrize(
-    "unit_diag", [True, False], ids=["unit_diag=True", "unit_diag=False"]
-)
-@pytest.mark.parametrize("is_complex", [True, False], ids=["complex", "real"])
-@pytest.mark.filterwarnings(
-    'ignore:Cannot cache compiled function "numba_funcified_fgraph"'
-)
-def test_solve_triangular(b_shape: tuple[int], lower, trans, unit_diag, is_complex):
-    if is_complex:
-        # TODO: Complex raises ValueError: To change to a dtype of a different size, the last axis must be contiguous,
-        #  why?
-        pytest.skip("Complex inputs currently not supported to solve_triangular")
-
-    complex_dtype = "complex64" if floatX.endswith("32") else "complex128"
-    dtype = complex_dtype if is_complex else floatX
-
-    A = pt.matrix("A", dtype=dtype)
-    b = pt.tensor("b", shape=b_shape, dtype=dtype)
-
-    def A_func(x):
-        x = x @ x.conj().T
-        x_tri = pt.linalg.cholesky(x, lower=lower).astype(dtype)
-
-        if unit_diag:
-            x_tri = pt.fill_diagonal(x_tri, 1.0)
-
-        return x_tri
-
-    solve_op = partial(
-        pt.linalg.solve_triangular, lower=lower, trans=trans, unit_diagonal=unit_diag
-    )
-
-    X = solve_op(A_func(A), b)
-    f = pytensor.function([A, b], X, mode="NUMBA")
-
-    A_val = np.random.normal(size=(5, 5))
-    b_val = np.random.normal(size=b_shape)
-
-    if is_complex:
-        A_val = A_val + np.random.normal(size=(5, 5)) * 1j
-        b_val = b_val + np.random.normal(size=b_shape) * 1j
-
-    X_np = f(A_val.copy(), b_val.copy())
-    A_val_transformed = transpose_func(A_func(A_val), trans).eval()
-    np.testing.assert_allclose(
-        A_val_transformed @ X_np,
-        b_val,
-        atol=1e-8 if floatX.endswith("64") else 1e-4,
-        rtol=1e-8 if floatX.endswith("64") else 1e-4,
-    )
-
-    compiled_fgraph = f.maker.fgraph
-    compare_numba_and_py(
-        compiled_fgraph.inputs,
-        compiled_fgraph.outputs,
-        [A_val, b_val],
-    )
-
-
-@pytest.mark.parametrize(
-    "lower, unit_diag, trans",
-    [(True, True, True), (False, False, False)],
-    ids=["lower_unit_trans", "defaults"],
-)
-def test_solve_triangular_grad(lower, unit_diag, trans):
-    A_val = np.random.normal(size=(5, 5)).astype(floatX)
-    b_val = np.random.normal(size=(5, 5)).astype(floatX)
-
-    # utt.verify_grad uses small perturbations to the input matrix to calculate the finite difference gradient. When
-    # a non-triangular matrix is passed to scipy.linalg.solve_triangular, no error is raise, but the result will be
-    # wrong, resulting in wrong gradients. As a result, it is necessary to add a mapping from the space of all matrices
-    # to the space of triangular matrices, and test the gradient of that entire graph.
-    def A_func_pt(x):
-        x = x @ x.conj().T
-        x_tri = pt.linalg.cholesky(x, lower=lower).astype(floatX)
-
-        if unit_diag:
-            n = A_val.shape[0]
-            x_tri = x_tri[np.diag_indices(n)].set(1.0)
-
-        return transpose_func(x_tri.astype(floatX), trans)
-
-    solve_op = partial(
-        pt.linalg.solve_triangular, lower=lower, trans=trans, unit_diagonal=unit_diag
-    )
-
-    utt.verify_grad(
-        lambda A, b: solve_op(A_func_pt(A), b),
-        [A_val.copy(), b_val.copy()],
-        mode="NUMBA",
-    )
-
-
-@pytest.mark.parametrize("overwrite_b", [True, False], ids=["inplace", "not_inplace"])
-def test_solve_triangular_overwrite_b_correct(overwrite_b):
-    # Regression test for issue #1233
-
-    rng = np.random.default_rng(utt.fetch_seed())
-    a_test_py = np.asfortranarray(rng.normal(size=(3, 3)))
-    a_test_py = np.tril(a_test_py)
-    b_test_py = np.asfortranarray(rng.normal(size=(3, 2)))
-
-    # .T.copy().T creates an f-contiguous copy of an f-contiguous array (otherwise the copy is c-contiguous)
-    a_test_nb = a_test_py.copy(order="F")
-    b_test_nb = b_test_py.copy(order="F")
-
-    op = SolveTriangular(
-        unit_diagonal=False,
-        lower=False,
-        check_finite=True,
-        b_ndim=2,
-        overwrite_b=overwrite_b,
-    )
-
-    a_pt = pt.matrix("a", shape=(3, 3))
-    b_pt = pt.matrix("b", shape=(3, 2))
-    out = op(a_pt, b_pt)
-
-    py_fn = pytensor.function([a_pt, b_pt], out, accept_inplace=True)
-    numba_fn = pytensor.function([a_pt, b_pt], out, accept_inplace=True, mode="NUMBA")
-
-    x_py = py_fn(a_test_py, b_test_py)
-    x_nb = numba_fn(a_test_nb, b_test_nb)
-
-    np.testing.assert_allclose(
-        py_fn(a_test_py, b_test_py), numba_fn(a_test_nb, b_test_nb)
-    )
-    np.testing.assert_allclose(b_test_py, b_test_nb)
-
-    if overwrite_b:
-        np.testing.assert_allclose(b_test_py, x_py)
-        np.testing.assert_allclose(b_test_nb, x_nb)
-
-
-@pytest.mark.parametrize("value", [np.nan, np.inf])
-@pytest.mark.filterwarnings(
-    'ignore:Cannot cache compiled function "numba_funcified_fgraph"'
-)
-def test_solve_triangular_raises_on_nan_inf(value):
-    A = pt.matrix("A")
-    b = pt.matrix("b")
-
-    X = pt.linalg.solve_triangular(A, b, check_finite=True)
-    f = pytensor.function([A, b], X, mode="NUMBA")
-    A_val = np.random.normal(size=(5, 5)).astype(floatX)
-    A_sym = A_val @ A_val.conj().T
-
-    A_tri = np.linalg.cholesky(A_sym).astype(floatX)
-    b = np.full((5, 1), value).astype(floatX)
-
-    with pytest.raises(
-        np.linalg.LinAlgError,
-        match=re.escape("Non-numeric values"),
-    ):
-        f(A_tri, b)
-
-
-@pytest.mark.parametrize("lower", [True, False], ids=["lower=True", "lower=False"])
-@pytest.mark.parametrize("trans", [True, False], ids=["trans=True", "trans=False"])
-def test_numba_Cholesky(lower, trans):
-    cov = pt.matrix("cov")
-
-    if trans:
-        cov_ = cov.T
-    else:
-        cov_ = cov
-    chol = pt.linalg.cholesky(cov_, lower=lower)
-
-    x = np.array([0.1, 0.2, 0.3]).astype(floatX)
-    val = np.eye(3).astype(floatX) + x[None, :] * x[:, None]
-
-    compare_numba_and_py([cov], [chol], [val])
-
-
-def test_numba_Cholesky_raises_on_nan_input():
-    test_value = rng.random(size=(3, 3)).astype(floatX)
-    test_value[0, 0] = np.nan
-
-    x = pt.tensor(dtype=floatX, shape=(3, 3))
-    x = x.T.dot(x)
-    g = pt.linalg.cholesky(x)
-    f = pytensor.function([x], g, mode="NUMBA")
-
-    with pytest.raises(np.linalg.LinAlgError, match=r"Non-numeric values"):
-        f(test_value)
-
-
-@pytest.mark.parametrize("on_error", ["nan", "raise"])
-def test_numba_Cholesky_raise_on(on_error):
-    test_value = rng.random(size=(3, 3)).astype(floatX)
-
-    x = pt.tensor(dtype=floatX, shape=(3, 3))
-    g = pt.linalg.cholesky(x, on_error=on_error)
-    f = pytensor.function([x], g, mode="NUMBA")
-
-    if on_error == "raise":
-        with pytest.raises(
-            np.linalg.LinAlgError, match=r"Input to cholesky is not positive definite"
-        ):
-            f(test_value)
-    else:
-        assert np.all(np.isnan(f(test_value)))
-
-
-@pytest.mark.parametrize("lower", [True, False], ids=["lower=True", "lower=False"])
-def test_numba_Cholesky_grad(lower):
-    rng = np.random.default_rng(utt.fetch_seed())
-    L = rng.normal(size=(5, 5)).astype(floatX)
-    X = L @ L.T
-
-    chol_op = partial(pt.linalg.cholesky, lower=lower)
-    utt.verify_grad(chol_op, [X], mode="NUMBA")
-
-
-def test_block_diag():
-    A = pt.matrix("A")
-    B = pt.matrix("B")
-    C = pt.matrix("C")
-    D = pt.matrix("D")
-    X = pt.linalg.block_diag(A, B, C, D)
-
-    A_val = np.random.normal(size=(5, 5)).astype(floatX)
-    B_val = np.random.normal(size=(3, 3)).astype(floatX)
-    C_val = np.random.normal(size=(2, 2)).astype(floatX)
-    D_val = np.random.normal(size=(4, 4)).astype(floatX)
-    compare_numba_and_py([A, B, C, D], [X], [A_val, B_val, C_val, D_val])
-
-
 def test_lamch():
     from scipy.linalg import get_lapack_funcs
 
@@ -328,171 +82,397 @@ def test_xgecon(ord_numba, ord_scipy):
     np.testing.assert_allclose(rcond, rcond2)
 
 
-@pytest.mark.parametrize("overwrite_a", [True, False])
-def test_getrf(overwrite_a):
-    from scipy.linalg import lu_factor
-
-    from pytensor.link.numba.dispatch.slinalg import _getrf
-
-    # TODO: Refactor this test to use compare_numba_and_py after we implement lu_factor in pytensor
-
-    @numba.njit()
-    def getrf(x, overwrite_a):
-        return _getrf(x, overwrite_a=overwrite_a)
-
-    x = np.random.normal(size=(5, 5)).astype(floatX)
-    x = np.asfortranarray(
-        x
-    )  # x needs to be fortran-contiguous going into getrf for the overwrite option to work
-
-    lu, ipiv = lu_factor(x, overwrite_a=False)
-    LU, IPIV, info = getrf(x, overwrite_a=overwrite_a)
-
-    assert info == 0
-    assert_allclose(LU, lu)
-
-    if overwrite_a:
-        assert_allclose(x, LU)
-
-    # TODO: It seems IPIV is 1-indexed in FORTRAN, so we need to subtract 1. I can't find evidence that scipy is doing
-    #  this, though.
-    assert_allclose(IPIV - 1, ipiv)
-
-
-@pytest.mark.parametrize("trans", [0, 1])
-@pytest.mark.parametrize("overwrite_a", [True, False])
-@pytest.mark.parametrize("overwrite_b", [True, False])
-@pytest.mark.parametrize("b_shape", [(5,), (5, 3)], ids=["b_1d", "b_2d"])
-def test_getrs(trans, overwrite_a, overwrite_b, b_shape):
-    from scipy.linalg import lu_factor
-    from scipy.linalg import lu_solve as sp_lu_solve
-
-    from pytensor.link.numba.dispatch.slinalg import _getrf, _getrs
-
-    # TODO: Refactor this test to use compare_numba_and_py after we implement lu_solve in pytensor
-
-    @numba.njit()
-    def lu_solve(a, b, trans, overwrite_a, overwrite_b):
-        lu, ipiv, info = _getrf(a, overwrite_a=overwrite_a)
-        x, info = _getrs(lu, b, ipiv, trans=trans, overwrite_b=overwrite_b)
-        return x, lu, info
-
-    a = np.random.normal(size=(5, 5)).astype(floatX)
-    b = np.random.normal(size=b_shape).astype(floatX)
-
-    # inputs need to be fortran-contiguous going into getrf and getrs for the overwrite option to work
-    a = np.asfortranarray(a)
-    b = np.asfortranarray(b)
-
-    lu_and_piv = lu_factor(a, overwrite_a=False)
-    x_sp = sp_lu_solve(lu_and_piv, b, trans, overwrite_b=False)
-
-    x, lu, info = lu_solve(
-        a, b, trans, overwrite_a=overwrite_a, overwrite_b=overwrite_b
+class TestSolves:
+    @pytest.mark.parametrize("lower", [True, False], ids=lambda x: f"lower={x}")
+    @pytest.mark.parametrize(
+        "overwrite_a, overwrite_b",
+        [(False, False), (True, False), (False, True)],
+        ids=["no_overwrite", "overwrite_a", "overwrite_b"],
     )
-    assert info == 0
-    if overwrite_a:
-        assert_allclose(a, lu)
-    if overwrite_b:
-        assert_allclose(b, x)
-
-    assert_allclose(x, x_sp)
-
-
-@pytest.mark.parametrize(
-    "b_shape",
-    [(5, 1), (5, 5), (5,)],
-    ids=["b_col_vec", "b_matrix", "b_vec"],
-)
-@pytest.mark.parametrize("assume_a", ["gen", "sym", "pos"], ids=str)
-@pytest.mark.filterwarnings(
-    'ignore:Cannot cache compiled function "numba_funcified_fgraph"'
-)
-def test_solve(b_shape: tuple[int], assume_a: Literal["gen", "sym", "pos"]):
-    A = pt.matrix("A", dtype=floatX)
-    b = pt.tensor("b", shape=b_shape, dtype=floatX)
-
-    A_val = np.asfortranarray(np.random.normal(size=(5, 5)).astype(floatX))
-    b_val = np.asfortranarray(np.random.normal(size=b_shape).astype(floatX))
-
-    def A_func(x):
-        if assume_a == "pos":
-            x = x @ x.T
-        elif assume_a == "sym":
-            x = (x + x.T) / 2
-        return x
-
-    X = pt.linalg.solve(
-        A_func(A),
-        b,
-        assume_a=assume_a,
-        b_ndim=len(b_shape),
+    @pytest.mark.parametrize(
+        "b_shape",
+        [(5, 1), (5, 5), (5,)],
+        ids=["b_col_vec", "b_matrix", "b_vec"],
+    )
+    @pytest.mark.parametrize("assume_a", ["gen", "sym", "pos"], ids=str)
+    def test_solve(
+        self,
+        b_shape: tuple[int],
+        assume_a: Literal["gen", "sym", "pos"],
+        lower: bool,
+        overwrite_a: bool,
+        overwrite_b: bool,
+    ):
+        if assume_a not in ("sym", "her", "pos") and not lower:
+            # Avoid redundant tests with lower=True and lower=False for non symmetric matrices
+            pytest.skip("Skipping redundant test already covered by lower=True")
+
+        def A_func(x):
+            if assume_a == "pos":
+                x = x @ x.T
+                x = np.tril(x) if lower else np.triu(x)
+            elif assume_a == "sym":
+                x = (x + x.T) / 2
+                n = x.shape[0]
+                # We have to set the unused triangle to something other than zero
+                # to see lapack destroying it.
+                x[np.triu_indices(n, 1) if lower else np.tril_indices(n, 1)] = np.pi
+            return x
+
+        A = pt.matrix("A", dtype=floatX)
+        b = pt.tensor("b", shape=b_shape, dtype=floatX)
+
+        rng = np.random.default_rng(418)
+        A_val = A_func(rng.normal(size=(5, 5))).astype(floatX)
+        b_val = rng.normal(size=b_shape).astype(floatX)
+
+        X = pt.linalg.solve(
+            A,
+            b,
+            assume_a=assume_a,
+            b_ndim=len(b_shape),
+        )
+
+        f, res = compare_numba_and_py(
+            [In(A, mutable=overwrite_a), In(b, mutable=overwrite_b)],
+            X,
+            test_inputs=[A_val, b_val],
+            inplace=True,
+            numba_mode=numba_inplace_mode,
+        )
+
+        op = f.maker.fgraph.outputs[0].owner.op
+        assert isinstance(op, Solve)
+        destroy_map = op.destroy_map
+        if overwrite_a and overwrite_b:
+            raise NotImplementedError(
+                "Test not implemented for simultaneous overwrite_a and overwrite_b, as that's not currently supported by PyTensor"
+            )
+        elif overwrite_a:
+            assert destroy_map == {0: [0]}
+        elif overwrite_b:
+            assert destroy_map == {0: [1]}
+        else:
+            assert destroy_map == {}
+
+        # Test with F_contiguous inputs
+        A_val_f_contig = np.copy(A_val, order="F")
+        b_val_f_contig = np.copy(b_val, order="F")
+        res_f_contig = f(A_val_f_contig, b_val_f_contig)
+        np.testing.assert_allclose(res_f_contig, res)
+        # Should always be destroyable
+        assert (A_val == A_val_f_contig).all() == (not overwrite_a)
+        assert (b_val == b_val_f_contig).all() == (not overwrite_b)
+
+        # Test with C_contiguous inputs
+        A_val_c_contig = np.copy(A_val, order="C")
+        b_val_c_contig = np.copy(b_val, order="C")
+        res_c_contig = f(A_val_c_contig, b_val_c_contig)
+        np.testing.assert_allclose(res_c_contig, res)
+        # We can destroy C-contiguous A arrays by inverting `tranpose/lower` at runtime
+        assert np.allclose(A_val_c_contig, A_val) == (not overwrite_a)
+        # b vectors are always f_contiguous if also c_contiguous
+        assert np.allclose(b_val_c_contig, b_val) == (
+            not (overwrite_b and b_val_c_contig.flags.f_contiguous)
+        )
+
+        # Test right results if inputs are not contiguous in either format
+        A_val_not_contig = np.repeat(A_val, 2, axis=0)[::2]
+        b_val_not_contig = np.repeat(b_val, 2, axis=0)[::2]
+        res_not_contig = f(A_val_not_contig, b_val_not_contig)
+        np.testing.assert_allclose(res_not_contig, res)
+        # Can never destroy non-contiguous inputs
+        np.testing.assert_allclose(A_val_not_contig, A_val)
+        np.testing.assert_allclose(b_val_not_contig, b_val)
+
+    @pytest.mark.parametrize("lower", [True, False], ids=lambda x: f"lower={x}")
+    @pytest.mark.parametrize(
+        "transposed", [False, True], ids=lambda x: f"transposed={x}"
     )
-    f = pytensor.function(
-        [pytensor.In(A, mutable=True), pytensor.In(b, mutable=True)], X, mode="NUMBA"
+    @pytest.mark.parametrize(
+        "overwrite_b", [False, True], ids=["no_overwrite", "overwrite_b"]
     )
-    op = f.maker.fgraph.outputs[0].owner.op
-
-    compare_numba_and_py([A, b], [X], test_inputs=[A_val, b_val], inplace=True)
-
-    # Calling this is destructive and will rewrite b_val to be the answer. Store copies of the inputs first.
-    A_val_copy = A_val.copy()
-    b_val_copy = b_val.copy()
-
-    X_np = f(A_val, b_val)
-
-    # overwrite_b is preferred when both inputs can be destroyed
-    assert op.destroy_map == {0: [1]}
+    @pytest.mark.parametrize(
+        "unit_diagonal", [True, False], ids=lambda x: f"unit_diagonal={x}"
+    )
+    @pytest.mark.parametrize(
+        "b_shape",
+        [(5, 1), (5, 5), (5,)],
+        ids=["b_col_vec", "b_matrix", "b_vec"],
+    )
+    @pytest.mark.parametrize("is_complex", [True, False], ids=["complex", "real"])
+    def test_solve_triangular(
+        self,
+        b_shape: tuple[int],
+        lower: bool,
+        transposed: bool,
+        unit_diagonal: bool,
+        is_complex: bool,
+        overwrite_b: bool,
+    ):
+        if is_complex:
+            # TODO: Complex raises ValueError: To change to a dtype of a different size, the last axis must be contiguous,
+            #  why?
+            pytest.skip("Complex inputs currently not supported to solve_triangular")
+
+        def A_func(x):
+            complex_dtype = "complex64" if floatX.endswith("32") else "complex128"
+            dtype = complex_dtype if is_complex else floatX
+
+            x = x @ x.conj().T
+            x_tri = scipy.linalg.cholesky(x, lower=lower).astype(dtype)
+
+            if unit_diagonal:
+                x_tri[np.diag_indices(x_tri.shape[0])] = 1.0
+
+            return x_tri
+
+        A = pt.matrix("A", dtype=floatX)
+        b = pt.tensor("b", shape=b_shape, dtype=floatX)
+
+        rng = np.random.default_rng(418)
+        A_val = A_func(rng.normal(size=(5, 5))).astype(floatX)
+        b_val = rng.normal(size=b_shape).astype(floatX)
+
+        X = pt.linalg.solve_triangular(
+            A,
+            b,
+            lower=lower,
+            trans="N" if (not transposed) else ("C" if is_complex else "T"),
+            unit_diagonal=unit_diagonal,
+            b_ndim=len(b_shape),
+        )
+
+        f, res = compare_numba_and_py(
+            [A, In(b, mutable=overwrite_b)],
+            X,
+            test_inputs=[A_val, b_val],
+            inplace=True,
+            numba_mode=numba_inplace_mode,
+        )
+
+        op = f.maker.fgraph.outputs[0].owner.op
+        assert isinstance(op, SolveTriangular)
+        destroy_map = op.destroy_map
+        if overwrite_b:
+            assert destroy_map == {0: [1]}
+        else:
+            assert destroy_map == {}
+
+        # Test with F_contiguous inputs
+        A_val_f_contig = np.copy(A_val, order="F")
+        b_val_f_contig = np.copy(b_val, order="F")
+        res_f_contig = f(A_val_f_contig, b_val_f_contig)
+        np.testing.assert_allclose(res_f_contig, res)
+        # solve_triangular never destroys A
+        np.testing.assert_allclose(A_val, A_val_f_contig)
+        # b Should always be destroyable
+        assert (b_val == b_val_f_contig).all() == (not overwrite_b)
+
+        # Test with C_contiguous inputs
+        A_val_c_contig = np.copy(A_val, order="C")
+        b_val_c_contig = np.copy(b_val, order="C")
+        res_c_contig = f(A_val_c_contig, b_val_c_contig)
+        np.testing.assert_allclose(res_c_contig, res)
+        np.testing.assert_allclose(A_val_c_contig, A_val)
+        # b c_contiguous vectors are also f_contiguous and destroyable
+        assert np.allclose(b_val_c_contig, b_val) == (
+            not (overwrite_b and b_val_c_contig.flags.f_contiguous)
+        )
+
+        # Test with non-contiguous inputs
+        A_val_not_contig = np.repeat(A_val, 2, axis=0)[::2]
+        b_val_not_contig = np.repeat(b_val, 2, axis=0)[::2]
+        res_not_contig = f(A_val_not_contig, b_val_not_contig)
+        np.testing.assert_allclose(res_not_contig, res)
+        np.testing.assert_allclose(A_val_not_contig, A_val)
+        # Can never destroy non-contiguous inputs
+        np.testing.assert_allclose(b_val_not_contig, b_val)
+
+    @pytest.mark.parametrize("value", [np.nan, np.inf])
+    def test_solve_triangular_raises_on_nan_inf(self, value):
+        A = pt.matrix("A")
+        b = pt.matrix("b")
+
+        X = pt.linalg.solve_triangular(A, b, check_finite=True)
+        f = pytensor.function([A, b], X, mode="NUMBA")
+        A_val = np.random.normal(size=(5, 5)).astype(floatX)
+        A_sym = A_val @ A_val.conj().T
+
+        A_tri = np.linalg.cholesky(A_sym).astype(floatX)
+        b = np.full((5, 1), value).astype(floatX)
 
-    # Confirm inputs were destroyed by checking against the copies
-    assert (A_val == A_val_copy).all() == (op.destroy_map.get(0, None) != [0])
-    assert (b_val == b_val_copy).all() == (op.destroy_map.get(0, None) != [1])
+        with pytest.raises(
+            np.linalg.LinAlgError,
+            match=re.escape("Non-numeric values"),
+        ):
+            f(A_tri, b)
 
-    ATOL = 1e-8 if floatX.endswith("64") else 1e-4
-    RTOL = 1e-8 if floatX.endswith("64") else 1e-4
+    @pytest.mark.parametrize("lower", [True, False], ids=lambda x: f"lower = {x}")
+    @pytest.mark.parametrize(
+        "overwrite_b", [False, True], ids=["no_overwrite", "overwrite_b"]
+    )
+    @pytest.mark.parametrize(
+        "b_func, b_shape",
+        [(pt.matrix, (5, 1)), (pt.matrix, (5, 5)), (pt.vector, (5,))],
+        ids=["b_col_vec", "b_matrix", "b_vec"],
+    )
+    def test_cho_solve(
+        self, b_func, b_shape: tuple[int, ...], lower: bool, overwrite_b: bool
+    ):
+        def A_func(x):
+            x = x @ x.conj().T
+            x = scipy.linalg.cholesky(x, lower=lower)
+            return x
+
+        A = pt.matrix("A", dtype=floatX)
+        b = pt.tensor("b", shape=b_shape, dtype=floatX)
+
+        rng = np.random.default_rng(418)
+        A_val = A_func(rng.normal(size=(5, 5))).astype(floatX)
+        b_val = rng.normal(size=b_shape).astype(floatX)
+
+        X = pt.linalg.cho_solve(
+            (A, lower),
+            b,
+            b_ndim=len(b_shape),
+        )
+
+        f, res = compare_numba_and_py(
+            [A, In(b, mutable=overwrite_b)],
+            X,
+            test_inputs=[A_val, b_val],
+            inplace=True,
+            numba_mode=numba_inplace_mode,
+        )
+
+        op = f.maker.fgraph.outputs[0].owner.op
+        assert isinstance(op, CholeskySolve)
+        destroy_map = op.destroy_map
+        if overwrite_b:
+            assert destroy_map == {0: [1]}
+        else:
+            assert destroy_map == {}
+
+        # Test with F_contiguous inputs
+        A_val_f_contig = np.copy(A_val, order="F")
+        b_val_f_contig = np.copy(b_val, order="F")
+        res_f_contig = f(A_val_f_contig, b_val_f_contig)
+        np.testing.assert_allclose(res_f_contig, res)
+        # cho_solve never destroys A
+        np.testing.assert_allclose(A_val, A_val_f_contig)
+        # b Should always be destroyable
+        assert (b_val == b_val_f_contig).all() == (not overwrite_b)
+
+        # Test with C_contiguous inputs
+        A_val_c_contig = np.copy(A_val, order="C")
+        b_val_c_contig = np.copy(b_val, order="C")
+        res_c_contig = f(A_val_c_contig, b_val_c_contig)
+        np.testing.assert_allclose(res_c_contig, res)
+        np.testing.assert_allclose(A_val_c_contig, A_val)
+        # b c_contiguous vectors are also f_contiguous and destroyable
+        assert np.allclose(b_val_c_contig, b_val) == (
+            not (overwrite_b and b_val_c_contig.flags.f_contiguous)
+        )
+
+        # Test with non-contiguous inputs
+        A_val_not_contig = np.repeat(A_val, 2, axis=0)[::2]
+        b_val_not_contig = np.repeat(b_val, 2, axis=0)[::2]
+        res_not_contig = f(A_val_not_contig, b_val_not_contig)
+        np.testing.assert_allclose(res_not_contig, res)
+        np.testing.assert_allclose(A_val_not_contig, A_val)
+        # Can never destroy non-contiguous inputs
+        np.testing.assert_allclose(b_val_not_contig, b_val)
+
+
+@pytest.mark.parametrize("lower", [True, False], ids=lambda x: f"lower={x}")
+@pytest.mark.parametrize(
+    "overwrite_a", [False, True], ids=["no_overwrite", "overwrite_a"]
+)
+def test_cholesky(lower: bool, overwrite_a: bool):
+    cov = pt.matrix("cov")
+    chol = pt.linalg.cholesky(cov, lower=lower)
 
-    # Confirm b_val is used to store to solution
-    np.testing.assert_allclose(X_np, b_val, atol=ATOL, rtol=RTOL)
-    assert not np.allclose(b_val, b_val_copy)
+    x = np.array([0.1, 0.2, 0.3]).astype(floatX)
+    val = np.eye(3).astype(floatX) + x[None, :] * x[:, None]
 
-    # Test that the result is numerically correct. Need to use the unmodified copy
-    np.testing.assert_allclose(
-        A_func(A_val_copy) @ X_np, b_val_copy, atol=ATOL, rtol=RTOL
+    fn, res = compare_numba_and_py(
+        [In(cov, mutable=overwrite_a)],
+        [chol],
+        [val],
+        numba_mode=numba_inplace_mode,
+        inplace=True,
     )
 
-    # See the note in tensor/test_slinalg.py::test_solve_correctness for details about the setup here
-    utt.verify_grad(
-        lambda A, b: pt.linalg.solve(
-            A_func(A), b, lower=False, assume_a=assume_a, b_ndim=len(b_shape)
-        ),
-        [A_val_copy, b_val_copy],
-        mode="NUMBA",
-    )
+    op = fn.maker.fgraph.outputs[0].owner.op
+    assert isinstance(op, Cholesky)
+    destroy_map = op.destroy_map
+    if overwrite_a:
+        assert destroy_map == {0: [0]}
+    else:
+        assert destroy_map == {}
+
+    # Test F-contiguous input
+    val_f_contig = np.copy(val, order="F")
+    res_f_contig = fn(val_f_contig)
+    np.testing.assert_allclose(res_f_contig, res)
+    # Should always be destroyable
+    assert (val == val_f_contig).all() == (not overwrite_a)
+
+    # Test C-contiguous input
+    val_c_contig = np.copy(val, order="C")
+    res_c_contig = fn(val_c_contig)
+    np.testing.assert_allclose(res_c_contig, res)
+    # Cannot destroy C-contiguous input
+    np.testing.assert_allclose(val_c_contig, val)
+
+    # Test non-contiguous input
+    val_not_contig = np.repeat(val, 2, axis=0)[::2]
+    res_not_contig = fn(val_not_contig)
+    np.testing.assert_allclose(res_not_contig, res)
+    # Cannot destroy non-contiguous input
+    np.testing.assert_allclose(val_not_contig, val)
+
+
+def test_cholesky_raises_on_nan_input():
+    test_value = rng.random(size=(3, 3)).astype(floatX)
+    test_value[0, 0] = np.nan
 
+    x = pt.tensor(dtype=floatX, shape=(3, 3))
+    x = x.T.dot(x)
+    g = pt.linalg.cholesky(x)
+    f = pytensor.function([x], g, mode="NUMBA")
 
-@pytest.mark.parametrize(
-    "b_func, b_size",
-    [(pt.matrix, (5, 1)), (pt.matrix, (5, 5)), (pt.vector, (5,))],
-    ids=["b_col_vec", "b_matrix", "b_vec"],
-)
-@pytest.mark.parametrize("lower", [True, False], ids=lambda x: f"lower = {x}")
-def test_cho_solve(b_func, b_size, lower):
-    A = pt.matrix("A", dtype=floatX)
-    b = b_func("b", dtype=floatX)
+    with pytest.raises(np.linalg.LinAlgError, match=r"Non-numeric values"):
+        f(test_value)
 
-    C = pt.linalg.cholesky(A, lower=lower)
-    X = pt.linalg.cho_solve((C, lower), b)
-    f = pytensor.function([A, b], X, mode="NUMBA")
 
-    A = np.random.normal(size=(5, 5)).astype(floatX)
-    A = A @ A.conj().T
+@pytest.mark.parametrize("on_error", ["nan", "raise"])
+def test_cholesky_raise_on(on_error):
+    test_value = rng.random(size=(3, 3)).astype(floatX)
+
+    x = pt.tensor(dtype=floatX, shape=(3, 3))
+    g = pt.linalg.cholesky(x, on_error=on_error)
+    f = pytensor.function([x], g, mode="NUMBA")
 
-    b = np.random.normal(size=b_size)
-    b = b.astype(floatX)
+    if on_error == "raise":
+        with pytest.raises(
+            np.linalg.LinAlgError, match=r"Input to cholesky is not positive definite"
+        ):
+            f(test_value)
+    else:
+        assert np.all(np.isnan(f(test_value)))
 
-    X_np = f(A, b)
 
-    ATOL = 1e-8 if floatX.endswith("64") else 1e-4
-    RTOL = 1e-8 if floatX.endswith("64") else 1e-4
+def test_block_diag():
+    A = pt.matrix("A")
+    B = pt.matrix("B")
+    C = pt.matrix("C")
+    D = pt.matrix("D")
+    X = pt.linalg.block_diag(A, B, C, D)
 
-    np.testing.assert_allclose(A @ X_np, b, atol=ATOL, rtol=RTOL)
+    A_val = np.random.normal(size=(5, 5)).astype(floatX)
+    B_val = np.random.normal(size=(3, 3)).astype(floatX)
+    C_val = np.random.normal(size=(2, 2)).astype(floatX)
+    D_val = np.random.normal(size=(4, 4)).astype(floatX)
+    compare_numba_and_py([A, B, C, D], [X], [A_val, B_val, C_val, D_val])
diff --git a/tests/tensor/test_basic.py b/tests/tensor/test_basic.py
index 9be5044f9..dee0023ef 100644
--- a/tests/tensor/test_basic.py
+++ b/tests/tensor/test_basic.py
@@ -2861,6 +2861,13 @@ class TestARange:
             assert np.all(f(2) == len(np.arange(0, 2)))
             assert np.all(f(0) == len(np.arange(0, 0)))
 
+    def test_static_shape(self):
+        assert np.arange(1, 10).shape == arange(1, 10).type.shape
+        assert np.arange(10, 1, -1).shape == arange(10, 1, -1).type.shape
+        assert np.arange(1, -9, 2).shape == arange(1, -9, 2).type.shape
+        assert np.arange(1.3, 17.48, 2.67).shape == arange(1.3, 17.48, 2.67).type.shape
+        assert np.arange(-64, 64).shape == arange(-64, 64).type.shape
+
 
 class TestNdGrid:
     def setup_method(self):
diff --git a/tests/tensor/test_extra_ops.py b/tests/tensor/test_extra_ops.py
index b03d59170..e8900ce5d 100644
--- a/tests/tensor/test_extra_ops.py
+++ b/tests/tensor/test_extra_ops.py
@@ -965,8 +965,10 @@ class TestUnravelIndex(utt.InferShapeTester):
             f_array_array = fn(indices, shape_array)
             np.testing.assert_equal(ref, f_array_array())
 
-            # shape given as an PyTensor variable
-            shape_symb = pytensor.shared(shape_array)
+            # shape given as a shared PyTensor variable with static shape
+            shape_symb = pytensor.shared(
+                shape_array, shape=shape_array.shape, strict=True
+            )
             f_array_symb = fn(indices, shape_symb)
             np.testing.assert_equal(ref, f_array_symb())
 
diff --git a/tests/tensor/test_math_scipy.py b/tests/tensor/test_math_scipy.py
index 8f7095020..e15293e97 100644
--- a/tests/tensor/test_math_scipy.py
+++ b/tests/tensor/test_math_scipy.py
@@ -18,7 +18,7 @@ from pytensor import function, grad
 from pytensor import tensor as pt
 from pytensor.compile.mode import get_default_mode
 from pytensor.configdefaults import config
-from pytensor.tensor import gammaincc, inplace, kv, kve, vector
+from pytensor.tensor import gammaincc, inplace, kn, kv, kve, vector
 from tests import unittest_tools as utt
 from tests.tensor.utils import (
     _good_broadcast_unary_chi2sf,
@@ -1220,3 +1220,17 @@ def test_kv():
         out.eval({v: test_v, x: test_x}),
         scipy.special.kv(test_v[:, None], test_x[None, :]),
     )
+
+
+def test_kn():
+    n = vector("n")
+    x = vector("x")
+
+    out = kn(n[:, None], x[None, :])
+    test_n = np.array([-3, 4, 0, 5], dtype=n.type.dtype)
+    test_x = np.linspace(0, 512, 10, dtype=x.type.dtype)
+
+    np.testing.assert_allclose(
+        out.eval({n: test_n, x: test_x}),
+        scipy.special.kn(test_n[:, None], test_x[None, :]),
+    )
diff --git a/tests/tensor/test_nlinalg.py b/tests/tensor/test_nlinalg.py
index 4b83446c5..c8ae3ac4c 100644
--- a/tests/tensor/test_nlinalg.py
+++ b/tests/tensor/test_nlinalg.py
@@ -152,6 +152,72 @@ def test_qr_modes():
         assert "name 'complete' is not defined" in str(e)
 
 
+@pytest.mark.parametrize(
+    "shape, gradient_test_case, mode",
+    (
+        [(s, c, "reduced") for s in [(3, 3), (6, 3), (3, 6)] for c in [0, 1, 2]]
+        + [(s, c, "complete") for s in [(3, 3), (6, 3), (3, 6)] for c in [0, 1, 2]]
+        + [(s, 0, "r") for s in [(3, 3), (6, 3), (3, 6)]]
+        + [((3, 3), 0, "raw")]
+    ),
+    ids=(
+        [
+            f"shape={s}, gradient_test_case={c}, mode=reduced"
+            for s in [(3, 3), (6, 3), (3, 6)]
+            for c in ["Q", "R", "both"]
+        ]
+        + [
+            f"shape={s}, gradient_test_case={c}, mode=complete"
+            for s in [(3, 3), (6, 3), (3, 6)]
+            for c in ["Q", "R", "both"]
+        ]
+        + [f"shape={s}, gradient_test_case=R, mode=r" for s in [(3, 3), (6, 3), (3, 6)]]
+        + ["shape=(3, 3), gradient_test_case=Q, mode=raw"]
+    ),
+)
+@pytest.mark.parametrize("is_complex", [True, False], ids=["complex", "real"])
+def test_qr_grad(shape, gradient_test_case, mode, is_complex):
+    rng = np.random.default_rng(utt.fetch_seed())
+
+    def _test_fn(x, case=2, mode="reduced"):
+        if case == 0:
+            return qr(x, mode=mode)[0].sum()
+        elif case == 1:
+            return qr(x, mode=mode)[1].sum()
+        elif case == 2:
+            Q, R = qr(x, mode=mode)
+            return Q.sum() + R.sum()
+
+    if is_complex:
+        pytest.xfail("Complex inputs currently not supported by verify_grad")
+
+    m, n = shape
+    a = rng.standard_normal(shape).astype(config.floatX)
+    if is_complex:
+        a += 1j * rng.standard_normal(shape).astype(config.floatX)
+
+    if mode == "raw":
+        with pytest.raises(NotImplementedError):
+            utt.verify_grad(
+                partial(_test_fn, case=gradient_test_case, mode=mode),
+                [a],
+                rng=np.random,
+            )
+
+    elif mode == "complete" and m > n:
+        with pytest.raises(AssertionError):
+            utt.verify_grad(
+                partial(_test_fn, case=gradient_test_case, mode=mode),
+                [a],
+                rng=np.random,
+            )
+
+    else:
+        utt.verify_grad(
+            partial(_test_fn, case=gradient_test_case, mode=mode), [a], rng=np.random
+        )
+
+
 class TestSvd(utt.InferShapeTester):
     op_class = SVD
 
diff --git a/tests/tensor/test_sharedvar.py b/tests/tensor/test_sharedvar.py
index b6cbbf7d1..436334b43 100644
--- a/tests/tensor/test_sharedvar.py
+++ b/tests/tensor/test_sharedvar.py
@@ -605,6 +605,7 @@ def makeSharedTester(
         def test_values_eq(self):
             # Test the type.values_eq[_approx] function
             dtype = self.dtype
+
             if dtype is None:
                 dtype = pytensor.config.floatX
 
@@ -691,9 +692,13 @@ def test_scalar_shared_deprecated():
 
 
 def test_get_vector_length():
-    x = pytensor.shared(np.array((2, 3, 4, 5)))
+    arr = np.array((2, 3, 4, 5))
+    x = pytensor.shared(arr, shape=arr.shape, strict=True)
     assert get_vector_length(x) == 4
 
+    with pytest.raises(ValueError):
+        get_vector_length(pytensor.shared(arr))
+
 
 def test_shared_masked_array_not_implemented():
     x = np.ma.masked_greater(np.array([1, 2, 3, 4]), 3)