Repository: pymc-devs/pytensor

Issue:
BUG: Error in shape for JAX tridiagonal solve inputs
### Describe the issue:

When calling `pt.linalg.solve` with `assume_a='tridiagonal'` on a problem A.dot(x) = b, the code fails because  `dl` and `du` should have the same length as `d`.

Additionally, if `b` is a vector, the code fails because `b` is supposed to have rank 2


### Reproducable code example:

```python
import pytensor
import pytensor.tensor as pt
import numpy as np

A = pt.matrix()
b = pt.vector()

tridiagonal_solve = pytensor.function([A, b],pt.linalg.solve(A, b,assume_a='tridiagonal'), mode='JAX')


N = 10
A = np.eye(N)
for i in range(N-1):
    A[i,i+1] = np.random.randn()
    A[i+1,i] = np.random.randn()

b = np.random.randn(N)

x_pt = tridiagonal_solve(A, b)
x_np = np.linalg.inv(A).dot(b)

print(x_pt, x_np)
assert np.all(np.isclose(x_pt, x_np))
```

### Error message:

```shell
<details>
Traceback (most recent call last):
  File "/Users/bfmaier/Dropbox/business/projects_and_clients/pymc/pytensor/pytensor/link/basic.py", line 665, in thunk
    outputs = fgraph_jit(*(x[0] for x in thunk_inputs))
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/var/folders/lk/6xn0tp_j2hz9nj1g893krwj80000gn/T/tmpem2rxrhz", line 3, in jax_funcified_fgraph
    tensor_variable_2 = solve(tensor_variable, tensor_variable_1)
                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/bfmaier/Dropbox/business/projects_and_clients/pymc/pytensor/pytensor/link/jax/dispatch/slinalg.py", line 57, in solve
    return jax.lax.linalg.tridiagonal_solve(dl, d, du, b, lower=lower)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: tridiagonal_solve() got an unexpected keyword argument 'lower'
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/Users/bfmaier/Dropbox/business/projects_and_clients/pymc/test_pytensor/test_jax.py", line 19, in <module>
    x_pt = tridiagonal_solve(A, b)
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/bfmaier/Dropbox/business/projects_and_clients/pymc/pytensor/pytensor/compile/function/types.py", line 1039, in __call__
    outputs = vm() if output_subset is None else vm(output_subset=output_subset)
              ^^^^
  File "/Users/bfmaier/Dropbox/business/projects_and_clients/pymc/pytensor/pytensor/link/basic.py", line 669, in thunk
    raise_with_op(self.fgraph, output_nodes[0], thunk)
  File "/Users/bfmaier/Dropbox/business/projects_and_clients/pymc/pytensor/pytensor/link/utils.py", line 526, in raise_with_op
    raise exc_value.with_traceback(exc_trace)
  File "/Users/bfmaier/Dropbox/business/projects_and_clients/pymc/pytensor/pytensor/link/basic.py", line 665, in thunk
    outputs = fgraph_jit(*(x[0] for x in thunk_inputs))
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/var/folders/lk/6xn0tp_j2hz9nj1g893krwj80000gn/T/tmpem2rxrhz", line 3, in jax_funcified_fgraph
    tensor_variable_2 = solve(tensor_variable, tensor_variable_1)
                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/bfmaier/Dropbox/business/projects_and_clients/pymc/pytensor/pytensor/link/jax/dispatch/slinalg.py", line 57, in solve
    return jax.lax.linalg.tridiagonal_solve(dl, d, du, b, lower=lower)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: tridiagonal_solve() got an unexpected keyword argument 'lower'
Apply node that caused the error: Solve{assume_a='tridiagonal', lower=False, check_finite=True, b_ndim=1, overwrite_a=False, overwrite_b=False}(<Matrix(float64, shape=(?, ?))>, <Vector(float64, shape=(?,))>)
Toposort index: 0
Inputs types: [TensorType(float64, shape=(None, None)), TensorType(float64, shape=(None,))]
Inputs shapes: [(10, 10), (10,)]
Inputs strides: [(80, 8), (8,)]
Inputs values: ['not shown', 'not shown']
Outputs clients: [[output[0](Solve{assume_a='tridiagonal', lower=False, check_finite=True, b_ndim=1, overwrite_a=False, overwrite_b=False}.0)]]

Backtrace when the node is created (use PyTensor flag traceback__limit=N to make it longer):
  File "/Users/bfmaier/Dropbox/business/projects_and_clients/pymc/test_pytensor/test_jax.py", line 8, in <module>
    tridiagonal_solve = pytensor.function([A, b],pt.linalg.solve(A, b,assume_a='tridiagonal'), mode='JAX')

HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.
</details>
```

### PyTensor version information:

<details>
floatX ({'float32', 'float16', 'float64'})
    Doc:  Default floating-point precision for python casts.

Note: float16 support is experimental, use at your own risk.
    Value:  float64

warn_float64 ({'pdb', 'raise', 'ignore', 'warn'})
    Doc:  Do an action when a tensor variable with float64 dtype is created.
    Value:  ignore

pickle_test_value (<bound method BoolParam._apply of <pytensor.configparser.BoolParam object at 0x100f775f0>>)
    Doc:  Dump test values while pickling model. If True, test values will be dumped with model.
    Value:  True

cast_policy ({'custom', 'numpy+floatX'})
    Doc:  Rules for implicit type casting
    Value:  custom

device (cpu)
    Doc:  Default device for computations. only cpu is supported for now
    Value:  cpu

conv__assert_shape (<bound method BoolParam._apply of <pytensor.configparser.BoolParam object at 0x101da5a30>>)
    Doc:  If True, AbstractConv* ops will verify that user-provided shapes match the runtime shapes (debugging option, may slow down compilation)
    Value:  False

print_global_stats (<bound method BoolParam._apply of <pytensor.configparser.BoolParam object at 0x101da5d90>>)
    Doc:  Print some global statistics (time spent) at the end
    Value:  False

unpickle_function (<bound method BoolParam._apply of <pytensor.configparser.BoolParam object at 0x101da5ac0>>)
    Doc:  Replace unpickled PyTensor functions with None. This is useful to unpickle old graphs that pickled them when it shouldn't
    Value:  True

<pytensor.configparser.ConfigParam object at 0x101da5a90>
    Doc:  Default compilation mode
    Value:  Mode

cxx (<class 'str'>)
    Doc:  The C++ compiler to use. Currently only g++ is supported, but supporting additional compilers should not be too difficult. If it is empty, no C++ code is compiled.
    Value:  /usr/bin/clang++

linker ({'c', 'py', 'cvm_nogc', 'vm', 'c|py', 'c|py_nogc', 'vm_nogc', 'cvm'})
    Doc:  Default linker used if the pytensor flags mode is Mode
    Value:  cvm

allow_gc (<bound method BoolParam._apply of <pytensor.configparser.BoolParam object at 0x101da5d30>>)
    Doc:  Do we default to delete intermediate results during PyTensor function calls? Doing so lowers the memory requirement, but asks that we reallocate memory at the next function call. This is implemented for the default linker, but may not work for all linkers.
    Value:  True

optimizer ({'o3', 'fast_compile', 'merge', 'None', 'unsafe', 'o2', 'o1', 'o4', 'fast_run'})
    Doc:  Default optimizer. If not None, will use this optimizer with the Mode
    Value:  o4

optimizer_verbose (<bound method BoolParam._apply of <pytensor.configparser.BoolParam object at 0x101da5a60>>)
    Doc:  Print information about rewrites that are applied during a graph transformation.
    Value:  False

optimizer_verbose_ignore (<class 'str'>)
    Doc:  Do not print information for rewrites with these names when `optimizer_verbose` is `True`. Separate names with ','
    Value:

on_opt_error ({'pdb', 'raise', 'ignore', 'warn'})
    Doc:  What to do when an optimization crashes: warn and skip it, raise the exception, or fall into the pdb debugger.
    Value:  warn

nocleanup (<bound method BoolParam._apply of <pytensor.configparser.BoolParam object at 0x101da56d0>>)
    Doc:  Suppress the deletion of code files that did not compile cleanly
    Value:  False

on_unused_input ({'raise', 'ignore', 'warn'})
    Doc:  What to do if a variable in the 'inputs' list of  pytensor.function() is not used in the graph.
    Value:  raise

gcc__cxxflags (<class 'str'>)
    Doc:  Extra compiler flags for gcc
    Value:

cmodule__warn_no_version (<bound method BoolParam._apply of <pytensor.configparser.BoolParam object at 0x101da6390>>)
    Doc:  If True, will print a warning when compiling one or more Op with C code that can't be cached because there is no c_code_cache_version() function associated to at least one of those Ops.
    Value:  False

cmodule__remove_gxx_opt (<bound method BoolParam._apply of <pytensor.configparser.BoolParam object at 0x101da59d0>>)
    Doc:  If True, will remove the -O* parameter passed to g++.This is useful to debug in gdb modules compiled by PyTensor.The parameter -g is passed by default to g++
    Value:  False

cmodule__compilation_warning (<bound method BoolParam._apply of <pytensor.configparser.BoolParam object at 0x101da5940>>)
    Doc:  If True, will print compilation warnings.
    Value:  False

cmodule__preload_cache (<bound method BoolParam._apply of <pytensor.configparser.BoolParam object at 0x101da5d60>>)
    Doc:  If set to True, will preload the C module cache at import time
    Value:  False

cmodule__age_thresh_use (<class 'int'>)
    Doc:  In seconds. The time after which PyTensor won't reuse a compile c module.
    Value:  2073600

cmodule__debug (<bound method BoolParam._apply of <pytensor.configparser.BoolParam object at 0x1011bae70>>)
    Doc:  If True, define a DEBUG macro (if not exists) for any compiled C code.
    Value:  False

compile__wait (<class 'int'>)
    Doc:  Time to wait before retrying to acquire the compile lock.
    Value:  5

compile__timeout (<class 'int'>)
    Doc:  In seconds, time that a process will wait before deciding to
    override an existing lock. An override only happens when the existing
    lock is held by the same owner *and* has not been 'refreshed' by this
    owner for more than this period. Refreshes are done every half timeout
    period for running processes.
    Value:  120

tensor__cmp_sloppy (<class 'int'>)
    Doc:  Relax pytensor.tensor.math._allclose (0) not at all, (1) a bit, (2) more
    Value:  0

lib__amdlibm (<bound method BoolParam._apply of <pytensor.configparser.BoolParam object at 0x101da5f70>>)
    Doc:  Use amd's amdlibm numerical library
    Value:  False

tensor__insert_inplace_optimizer_validate_nb (<class 'int'>)
    Doc:  -1: auto, if graph have less then 500 nodes 1, else 10
    Value:  -1

traceback__limit (<class 'int'>)
    Doc:  The number of stack to trace. -1 mean all.
    Value:  8

traceback__compile_limit (<class 'int'>)
    Doc:  The number of stack to trace to keep during compilation. -1 mean all. If greater then 0, will also make us save PyTensor internal stack trace.
    Value:  0

warn__ignore_bug_before ({'1.0.1', '0.8.2', '0.3', '0.9', '0.8', '1.0.4', '0.6', 'None', '0.4.1', '0.5', '0.8.1', '0.7', '1.0.5', '0.4', '1.0.3', '0.10', '1.0', 'all', '1.0.2'})
    Doc:  If 'None', we warn about all PyTensor bugs found by default. If 'all', we don't warn about PyTensor bugs found by default. If a version, we print only the warnings relative to PyTensor bugs found after that version. Warning for specific bugs can be configured with specific [warn] flags.
    Value:  0.9

exception_verbosity ({'low', 'high'})
    Doc:  If 'low', the text of exceptions will generally refer to apply nodes with short names such as Elemwise{add_no_inplace}. If 'high', some exceptions will also refer to apply nodes with long descriptions  like:
        A. Elemwise{add_no_inplace}
                B. log_likelihood_v_given_h
                C. log_likelihood_h
    Value:  low

print_test_value (<bound method BoolParam._apply of <pytensor.configparser.BoolParam object at 0x101da6090>>)
    Doc:  If 'True', the __eval__ of an PyTensor variable will return its test_value when this is available. This has the practical consequence that, e.g., in debugging `my_var` will print the same as `my_var.tag.test_value` when a test value is defined.
    Value:  False

compute_test_value ({'pdb', 'warn', 'raise', 'ignore', 'off'})
    Doc:  If 'True', PyTensor will run each op at graph build time, using Constants, SharedVariables and the tag 'test_value' as inputs to the function. This helps the user track down problems in the graph before it gets optimized.
    Value:  off

compute_test_value_opt ({'pdb', 'warn', 'raise', 'ignore', 'off'})
    Doc:  For debugging PyTensor optimization only. Same as compute_test_value, but is used during PyTensor optimization
    Value:  off

check_input (<bound method BoolParam._apply of <pytensor.configparser.BoolParam object at 0x101da6120>>)
    Doc:  Specify if types should check their input in their C code. It can be used to speed up compilation, reduce overhead (particularly for scalars) and reduce the number of generated C files.
    Value:  True

NanGuardMode__nan_is_error (<bound method BoolParam._apply of <pytensor.configparser.BoolParam object at 0x101da60f0>>)
    Doc:  Default value for nan_is_error
    Value:  True

NanGuardMode__inf_is_error (<bound method BoolParam._apply of <pytensor.configparser.BoolParam object at 0x101da6150>>)
    Doc:  Default value for inf_is_error
    Value:  True

NanGuardMode__big_is_error (<bound method BoolParam._apply of <pytensor.configparser.BoolParam object at 0x101da5fd0>>)
    Doc:  Default value for big_is_error
    Value:  True

NanGuardMode__action ({'pdb', 'raise', 'warn'})
    Doc:  What NanGuardMode does when it finds a problem
    Value:  raise

DebugMode__patience (<class 'int'>)
    Doc:  Optimize graph this many times to detect inconsistency
    Value:  10

DebugMode__check_c (<bound method BoolParam._apply of <pytensor.configparser.BoolParam object at 0x101da61e0>>)
    Doc:  Run C implementations where possible
    Value:  True

DebugMode__check_py (<bound method BoolParam._apply of <pytensor.configparser.BoolParam object at 0x101da6660>>)
    Doc:  Run Python implementations where possible
    Value:  True

DebugMode__check_finite (<bound method BoolParam._apply of <pytensor.configparser.BoolParam object at 0x101162d20>>)
    Doc:  True -> complain about NaN/Inf results
    Value:  True

DebugMode__check_strides (<class 'int'>)
    Doc:  Check that Python- and C-produced ndarrays have same strides. On difference: (0) - ignore, (1) warn, or (2) raise error
    Value:  0

DebugMode__warn_input_not_reused (<bound method BoolParam._apply of <pytensor.configparser.BoolParam object at 0x101da61b0>>)
    Doc:  Generate a warning when destroy_map or view_map says that an op works inplace, but the op did not reuse the input for its output.
    Value:  True

DebugMode__check_preallocated_output (<class 'str'>)
    Doc:  Test thunks with pre-allocated memory as output storage. This is a list of strings separated by ":". Valid values are: "initial" (initial storage in storage map, happens with Scan),"previous" (previously-returned memory), "c_contiguous", "f_contiguous", "strided" (positive and negative strides), "wrong_size" (larger and smaller dimensions), and "ALL" (all of the above).
    Value:

DebugMode__check_preallocated_output_ndim (<class 'int'>)
    Doc:  When testing with "strided" preallocated output memory, test all combinations of strides over that number of (inner-most) dimensions. You may want to reduce that number to reduce memory or time usage, but it is advised to keep a minimum of 2.
    Value:  4

profiling__time_thunks (<bound method BoolParam._apply of <pytensor.configparser.BoolParam object at 0x101da66f0>>)
    Doc:  Time individual thunks when profiling
    Value:  True

profiling__n_apply (<class 'int'>)
    Doc:  Number of Apply instances to print by default
    Value:  20

profiling__n_ops (<class 'int'>)
    Doc:  Number of Ops to print by default
    Value:  20

profiling__output_line_width (<class 'int'>)
    Doc:  Max line width for the profiling output
    Value:  512

profiling__min_memory_size (<class 'int'>)
    Doc:  For the memory profile, do not print Apply nodes if the size
                 of their outputs (in bytes) is lower than this threshold
    Value:  1024

profiling__min_peak_memory (<bound method BoolParam._apply of <pytensor.configparser.BoolParam object at 0x101da6600>>)
    Doc:  The min peak memory usage of the order
    Value:  False

profiling__destination (<class 'str'>)
    Doc:  File destination of the profiling output
    Value:  stderr

profiling__debugprint (<bound method BoolParam._apply of <pytensor.configparser.BoolParam object at 0x101da65a0>>)
    Doc:  Do a debugprint of the profiled functions
    Value:  False

profiling__ignore_first_call (<bound method BoolParam._apply of <pytensor.configparser.BoolParam object at 0x101da6570>>)
    Doc:  Do we ignore the first call of an PyTensor function.
    Value:  False

on_shape_error ({'raise', 'warn'})
    Doc:  warn: print a warning and use the default value. raise: raise an error
    Value:  warn

openmp (<bound method BoolParam._apply of <pytensor.configparser.BoolParam object at 0x101da6510>>)
    Doc:  Allow (or not) parallel computation on the CPU with OpenMP. This is the default value used when creating an Op that supports OpenMP parallelization. It is preferable to define it via the PyTensor configuration file ~/.pytensorrc or with the environment variable PYTENSOR_FLAGS. Parallelization is only done for some operations that implement it, and even for operations that implement parallelism, each operation is free to respect this flag or not. You can control the number of threads used with the environment variable OMP_NUM_THREADS. If it is set to 1, we disable openmp in PyTensor by default.
    Value:  False

openmp_elemwise_minsize (<class 'int'>)
    Doc:  If OpenMP is enabled, this is the minimum size of vectors for which the openmp parallelization is enabled in element wise ops.
    Value:  200000

optimizer_excluding (<class 'str'>)
    Doc:  When using the default mode, we will remove optimizer with these tags. Separate tags with ':'.
    Value:

optimizer_including (<class 'str'>)
    Doc:  When using the default mode, we will add optimizer with these tags. Separate tags with ':'.
    Value:

optimizer_requiring (<class 'str'>)
    Doc:  When using the default mode, we will require optimizer with these tags. Separate tags with ':'.
    Value:

optdb__position_cutoff (<class 'float'>)
    Doc:  Where to stop earlier during optimization. It represent the position of the optimizer where to stop.
    Value:  inf

optdb__max_use_ratio (<class 'float'>)
    Doc:  A ratio that prevent infinite loop in EquilibriumGraphRewriter.
    Value:  8.0

cycle_detection ({'fast', 'regular'})
    Doc:  If cycle_detection is set to regular, most inplaces are allowed,but it is slower. If cycle_detection is set to faster, less inplacesare allowed, but it makes the compilation faster.The interaction of which one give the lower peak memory usage iscomplicated and not predictable, so if you are close to the peakmemory usage, triyng both could give you a small gain.
    Value:  regular

check_stack_trace ({'log', 'raise', 'warn', 'off'})
    Doc:  A flag for checking the stack trace during the optimization process. default (off): does not check the stack trace of any optimization log: inserts a dummy stack trace that identifies the optimizationthat inserted the variable that had an empty stack trace.warn: prints a warning if a stack trace is missing and also a dummystack trace is inserted that indicates which optimization insertedthe variable that had an empty stack trace.raise: raises an exception if a stack trace is missing
    Value:  off

metaopt__verbose (<class 'int'>)
    Doc:  0 for silent, 1 for only warnings, 2 for full output withtimings and selected implementation
    Value:  0

unittests__rseed (<class 'str'>)
    Doc:  Seed to use for randomized unit tests. Special value 'random' means using a seed of None.
    Value:  666

warn__round (<bound method BoolParam._apply of <pytensor.configparser.BoolParam object at 0x101da6420>>)
    Doc:  Warn when using `tensor.round` with the default mode. Round changed its default from `half_away_from_zero` to `half_to_even` to have the same default as NumPy.
    Value:  False

profile (<bound method BoolParam._apply of <pytensor.configparser.BoolParam object at 0x101da6330>>)
    Doc:  If VM should collect profile information
    Value:  False

profile_optimizer (<bound method BoolParam._apply of <pytensor.configparser.BoolParam object at 0x101da6300>>)
    Doc:  If VM should collect optimizer profile information
    Value:  False

profile_memory (<bound method BoolParam._apply of <pytensor.configparser.BoolParam object at 0x101da62a0>>)
    Doc:  If VM should collect memory profile information and print it
    Value:  False

<pytensor.configparser.ConfigParam object at 0x101da6270>
    Doc:  Useful only for the VM Linkers. When lazy is None, auto detect if lazy evaluation is needed and use the appropriate version. If the C loop isn't being used and lazy is True, use the Stack VM; otherwise, use the Loop VM.
    Value:  None

numba__vectorize_target ({'cuda', 'parallel', 'cpu'})
    Doc:  Default target for numba.vectorize.
    Value:  cpu

numba__fastmath (<bound method BoolParam._apply of <pytensor.configparser.BoolParam object at 0x101da6780>>)
    Doc:  If True, use Numba's fastmath mode.
    Value:  True

numba__cache (<bound method BoolParam._apply of <pytensor.configparser.BoolParam object at 0x101da67b0>>)
    Doc:  If True, use Numba's file based caching.
    Value:  True

compiledir_format (<class 'str'>)
    Doc:  Format string for platform-dependent compiled module subdirectory
(relative to base_compiledir). Available keys: device, gxx_version,
hostname, numpy_version, platform, processor, pytensor_version,
python_bitwidth, python_int_bitwidth, python_version, short_platform.
Defaults to compiledir_%(short_platform)s-%(processor)s-
%(python_version)s-%(python_bitwidth)s.
    Value:  compiledir_%(short_platform)s-%(processor)s-%(python_version)s-%(python_bitwidth)s

<pytensor.configparser.ConfigParam object at 0x10153ea80>
    Doc:  platform-independent root directory for compiled modules
    Value:  /Users/bfmaier/.pytensor

<pytensor.configparser.ConfigParam object at 0x10104eae0>
    Doc:  platform-dependent cache directory for compiled modules
    Value:  /Users/bfmaier/.pytensor/compiledir_macOS-15.5-arm64-arm-64bit-arm-3.12.10-64

blas__ldflags (<class 'str'>)
    Doc:  lib[s] to include for [Fortran] level-3 blas implementation
    Value:  -framework Accelerate -Wl,-rpath,/Users/bfmaier/miniconda3/envs/pytensor-dev/lib

blas__check_openmp (<bound method BoolParam._apply of <pytensor.configparser.BoolParam object at 0x1021b24e0>>)
    Doc:  Check for openmp library conflict.
WARNING: Setting this to False leaves you open to wrong results in blas-related operations.
    Value:  True

scan__allow_gc (<bound method BoolParam._apply of <pytensor.configparser.BoolParam object at 0x115a79010>>)
    Doc:  Allow/disallow gc inside of Scan (default: False)
    Value:  False

scan__allow_output_prealloc (<bound method BoolParam._apply of <pytensor.configparser.BoolParam object at 0x125de8da0>>)
    Doc:  Allow/disallow memory preallocation for outputs inside of scan (default: True)
    Value:  True
</details>

### Context for the issue:

_No response_

Hints from issue comments:

Solution patch:
diff --git a/pytensor/link/jax/dispatch/slinalg.py b/pytensor/link/jax/dispatch/slinalg.py
index 3d6af0001..855052b12 100644
--- a/pytensor/link/jax/dispatch/slinalg.py
+++ b/pytensor/link/jax/dispatch/slinalg.py
@@ -46,6 +46,7 @@ def jax_funcify_Cholesky(op, **kwargs):
 def jax_funcify_Solve(op, **kwargs):
     assume_a = op.assume_a
     lower = op.lower
+    b_is_vec = op.b_ndim == 1
 
     if assume_a == "tridiagonal":
         # jax.scipy.solve does not yet support tridiagonal matrices
@@ -54,7 +55,20 @@ def jax_funcify_Solve(op, **kwargs):
             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)
+
+            # jax requires dl and du to have the same shape as d
+            dl = jax.numpy.pad(dl, (1, 0))
+            du = jax.numpy.pad(du, (0, 1))
+
+            if b_is_vec:
+                b = jax.numpy.expand_dims(b, -1)
+
+            res = jax.lax.linalg.tridiagonal_solve(dl, d, du, b)
+
+            if b_is_vec:
+                return jax.numpy.squeeze(res, -1)
+
+            return res
 
     else:
         if assume_a not in ("gen", "sym", "her", "pos"):
diff --git a/pytensor/link/numba/dispatch/basic.py b/pytensor/link/numba/dispatch/basic.py
index 4938ecc42..845d6afc7 100644
--- a/pytensor/link/numba/dispatch/basic.py
+++ b/pytensor/link/numba/dispatch/basic.py
@@ -31,7 +31,6 @@ from pytensor.link.utils import (
     fgraph_to_python,
 )
 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
@@ -466,7 +465,7 @@ def numba_funcify_ArgSortOp(op, node, **kwargs):
             axis = axis.item()
 
             Y = np.swapaxes(X, axis, 0)
-            result = np.empty_like(Y)
+            result = np.empty_like(Y, dtype="int64")
 
             indices = list(np.ndindex(Y.shape[1:]))
 
@@ -607,25 +606,6 @@ def numba_funcify_Dot(op, node, **kwargs):
     return dot
 
 
-@numba_funcify.register(Softplus)
-def numba_funcify_Softplus(op, node, **kwargs):
-    x_dtype = np.dtype(node.inputs[0].dtype)
-
-    @numba_njit
-    def softplus(x):
-        if x < -37.0:
-            value = np.exp(x)
-        elif x < 18.0:
-            value = np.log1p(np.exp(x))
-        elif x < 33.3:
-            value = x + np.exp(-x)
-        else:
-            value = x
-        return direct_cast(value, x_dtype)
-
-    return softplus
-
-
 @numba_funcify.register(Solve)
 def numba_funcify_Solve(op, node, **kwargs):
     assume_a = op.assume_a
@@ -689,11 +669,6 @@ def numba_funcify_BatchedDot(op, node, **kwargs):
     return batched_dot
 
 
-# NOTE: The remaining `pytensor.tensor.blas` `Op`s appear unnecessary, because
-# they're only used to optimize basic `Dot` nodes, and those GEMV and GEMM
-# optimizations are apparently already performed by Numba
-
-
 @numba_funcify.register(IfElse)
 def numba_funcify_IfElse(op, **kwargs):
     n_outs = op.n_outs
diff --git a/pytensor/link/numba/dispatch/elemwise.py b/pytensor/link/numba/dispatch/elemwise.py
index 9fd81dadc..7244762b9 100644
--- a/pytensor/link/numba/dispatch/elemwise.py
+++ b/pytensor/link/numba/dispatch/elemwise.py
@@ -561,7 +561,7 @@ def numba_funcify_Argmax(op, node, **kwargs):
 
         @numba_basic.numba_njit(inline="always")
         def argmax(x):
-            return 0
+            return np.array(0, dtype="int64")
 
     else:
         axes = tuple(int(ax) for ax in axis)
diff --git a/pytensor/link/numba/dispatch/linalg/decomposition/lu.py b/pytensor/link/numba/dispatch/linalg/decomposition/lu.py
index 570c024b0..739f0a699 100644
--- a/pytensor/link/numba/dispatch/linalg/decomposition/lu.py
+++ b/pytensor/link/numba/dispatch/linalg/decomposition/lu.py
@@ -30,7 +30,7 @@ def _lu_factor_to_lu(a, dtype, overwrite_a):
     # Fortran is 1 indexed, so we need to subtract 1 from the IPIV array
     IPIV = IPIV - 1
     p_inv = _pivot_to_permutation(IPIV, dtype=dtype)
-    perm = np.argsort(p_inv)
+    perm = np.argsort(p_inv).astype("int32")
 
     return perm, L, U
 
diff --git a/pytensor/link/numba/dispatch/nlinalg.py b/pytensor/link/numba/dispatch/nlinalg.py
index 860560d0a..3271b5bd2 100644
--- a/pytensor/link/numba/dispatch/nlinalg.py
+++ b/pytensor/link/numba/dispatch/nlinalg.py
@@ -52,7 +52,7 @@ def numba_funcify_Det(op, node, **kwargs):
 
     @numba_basic.numba_njit(inline="always")
     def det(x):
-        return numba_basic.direct_cast(np.linalg.det(inputs_cast(x)), out_dtype)
+        return np.array(np.linalg.det(inputs_cast(x))).astype(out_dtype)
 
     return det
 
@@ -68,8 +68,8 @@ def numba_funcify_SLogDet(op, node, **kwargs):
     def slogdet(x):
         sign, det = np.linalg.slogdet(inputs_cast(x))
         return (
-            numba_basic.direct_cast(sign, out_dtype_1),
-            numba_basic.direct_cast(det, out_dtype_2),
+            np.array(sign).astype(out_dtype_1),
+            np.array(det).astype(out_dtype_2),
         )
 
     return slogdet
diff --git a/pytensor/link/numba/dispatch/scalar.py b/pytensor/link/numba/dispatch/scalar.py
index e9b637b00..7e4703c8d 100644
--- a/pytensor/link/numba/dispatch/scalar.py
+++ b/pytensor/link/numba/dispatch/scalar.py
@@ -28,7 +28,7 @@ from pytensor.scalar.basic import (
     Second,
     Switch,
 )
-from pytensor.scalar.math import Erf, Erfc, GammaLn, Log1mexp, Sigmoid
+from pytensor.scalar.math import Erf, Erfc, GammaLn, Log1mexp, Sigmoid, Softplus
 
 
 @numba_funcify.register(ScalarOp)
@@ -312,3 +312,22 @@ def erfc(x):
 @numba_funcify.register(Erfc)
 def numba_funcify_Erfc(op, **kwargs):
     return numba_basic.global_numba_func(erfc)
+
+
+@numba_funcify.register(Softplus)
+def numba_funcify_Softplus(op, node, **kwargs):
+    out_dtype = np.dtype(node.outputs[0].type.dtype)
+
+    @numba_basic.numba_njit
+    def softplus(x):
+        if x < -37.0:
+            value = np.exp(x)
+        elif x < 18.0:
+            value = np.log1p(np.exp(x))
+        elif x < 33.3:
+            value = x + np.exp(-x)
+        else:
+            value = x
+        return numba_basic.direct_cast(value, out_dtype)
+
+    return softplus
diff --git a/pytensor/tensor/_linalg/solve/rewriting.py b/pytensor/tensor/_linalg/solve/rewriting.py
index ff1c74cde..9ea8db37f 100644
--- a/pytensor/tensor/_linalg/solve/rewriting.py
+++ b/pytensor/tensor/_linalg/solve/rewriting.py
@@ -14,16 +14,22 @@ from pytensor.tensor.slinalg import Solve, lu_factor, lu_solve
 from pytensor.tensor.variable import TensorVariable
 
 
-def decompose_A(A, assume_a):
+def decompose_A(A, assume_a, check_finite):
     if assume_a == "gen":
-        return lu_factor(A, check_finite=False)
+        return lu_factor(A, check_finite=check_finite)
     else:
         raise NotImplementedError
 
 
-def solve_lu_decomposed_system(A_decomp, b, b_ndim, assume_a, transposed=False):
-    if assume_a == "gen":
-        return lu_solve(A_decomp, b, b_ndim=b_ndim, trans=transposed)
+def solve_lu_decomposed_system(A_decomp, b, transposed=False, *, core_solve_op: Solve):
+    if core_solve_op.assume_a == "gen":
+        return lu_solve(
+            A_decomp,
+            b,
+            trans=transposed,
+            b_ndim=core_solve_op.b_ndim,
+            check_finite=core_solve_op.check_finite,
+        )
     else:
         raise NotImplementedError
 
@@ -102,14 +108,19 @@ def _split_lu_solve_steps(
         ):
             return None
 
-    A_decomp = decompose_A(A, assume_a=assume_a)
+    # If any Op had check_finite=True, we also do it for the LU decomposition
+    check_finite_decomp = False
+    for client, _ in A_solve_clients_and_transpose:
+        if client.op.core_op.check_finite:
+            check_finite_decomp = True
+            break
+    A_decomp = decompose_A(A, assume_a=assume_a, check_finite=check_finite_decomp)
 
     replacements = {}
     for client, transposed in A_solve_clients_and_transpose:
         _, b = client.inputs
-        b_ndim = client.op.core_op.b_ndim
         new_x = solve_lu_decomposed_system(
-            A_decomp, b, b_ndim=b_ndim, assume_a=assume_a, transposed=transposed
+            A_decomp, b, transposed=transposed, core_solve_op=client.op.core_op
         )
         [old_x] = client.outputs
         new_x = atleast_Nd(new_x, n=old_x.type.ndim).astype(old_x.type.dtype)
diff --git a/pytensor/tensor/blas.py b/pytensor/tensor/blas.py
index 0977b500a..fc8afcea5 100644
--- a/pytensor/tensor/blas.py
+++ b/pytensor/tensor/blas.py
@@ -113,23 +113,22 @@ from pytensor.tensor.type import DenseTensorType, tensor
 _logger = logging.getLogger("pytensor.tensor.blas")
 
 
-# If check_init_y() == True we need to initialize y when beta == 0.
-def check_init_y():
-    # TODO: What is going on here?
+def must_initialize_y_gemv():
+    # Check whether Scipy GEMV could output nan if y in not initialized
     from scipy.linalg.blas import get_blas_funcs
 
-    if check_init_y._result is None:
-        y = float("NaN") * np.ones((2,))
+    if must_initialize_y_gemv._result is None:
+        y = np.full((2,), np.nan)
         x = np.ones((2,))
         A = np.ones((2, 2))
         gemv = get_blas_funcs("gemv", dtype=y.dtype)
         gemv(1.0, A.T, x, 0.0, y, overwrite_y=True, trans=True)
-        check_init_y._result = np.isnan(y).any()
+        must_initialize_y_gemv._result = np.isnan(y).any()
 
-    return check_init_y._result
+    return must_initialize_y_gemv._result
 
 
-check_init_y._result = None  # type: ignore
+must_initialize_y_gemv._result = None  # type: ignore
 
 
 class Gemv(Op):
@@ -197,7 +196,13 @@ class Gemv(Op):
                     f"(beta * y + alpha * dot(A, x)). y: {y.shape}, A: {A.shape}, x: {x.shape}"
                 )
 
-            if beta == 0 and check_init_y():
+            if beta == 0 and must_initialize_y_gemv():
+                # Most BLAS implementations of GEMV ignore y=nan when beta=0
+                # PyTensor considers that the correct behavior,
+                # and even exploits it to avoid copying or initializing outputs.
+                # By deciding to exploit this, however, it becomes our responsibility
+                # to ensure the behavior even in the rare cases BLAS deviates,
+                # or users will get errors, even for graphs that had no nan to begin with.
                 y.fill(0)
 
             # Here I suppose that A is in c order. If we don't make it
diff --git a/pytensor/tensor/blas_c.py b/pytensor/tensor/blas_c.py
index 951316302..0ef0a1f47 100644
--- a/pytensor/tensor/blas_c.py
+++ b/pytensor/tensor/blas_c.py
@@ -336,11 +336,12 @@ cger_no_inplace = CGer(False)
 # ##### ####### #######
 
 
-def gemv_c_code(y, A, x, z, alpha, beta, fail, force_init_beta=False, params=None):
+def gemv_c_code(y, A, x, z, alpha, beta, fail, must_initialize_y=False, params=None):
     """
     z <- beta * y + alpha * dot(A, x)
 
     where A is a matrix, y and x are vectors (ergo z is vector)
+    z = y if inplace else y.copy()
     """
     code = """
 
@@ -400,17 +401,11 @@ def gemv_c_code(y, A, x, z, alpha, beta, fail, force_init_beta=False, params=Non
         }
         if (dbeta != 0)
         {
+            // If dbeta is zero, we avoid doing the copy
             if (PyArray_CopyInto(%(z)s, %(y)s) != 0) {
                 %(fail)s
             }
         }
-        else if (%(force_init_beta)d)
-        {
-            PyObject *zero = PyFloat_FromDouble(0.);
-            if (zero == NULL) %(fail)s;
-            if (PyArray_FillWithScalar(%(z)s, zero) != 0) %(fail)s;
-            Py_DECREF(zero);
-        }
     }
     else
     {
@@ -422,6 +417,20 @@ def gemv_c_code(y, A, x, z, alpha, beta, fail, force_init_beta=False, params=Non
         }
     }
 
+    if (%(must_initialize_y)d && dbeta == 0)
+    {
+        // Most BLAS implementations of GEMV ignore y=nan when beta=0
+        // PyTensor considers that the correct behavior,
+        // and even exploits it to avoid copying or initializing outputs.
+        // By deciding to exploit this, however, it becomes our responsibility
+        // to ensure the behavior even in the rare cases BLAS deviates,
+        // or users will get errors, even for graphs that had no nan to begin with.
+        PyObject *zero = PyFloat_FromDouble(0.);
+        if (zero == NULL) %(fail)s;
+        if (PyArray_FillWithScalar(%(z)s, zero) != 0) %(fail)s;
+        Py_DECREF(zero);
+    }
+
     {
         int NA0 = PyArray_DIMS(%(A)s)[0];
         int NA1 = PyArray_DIMS(%(A)s)[1];
@@ -491,13 +500,13 @@ def gemv_c_code(y, A, x, z, alpha, beta, fail, force_init_beta=False, params=Non
 
                 if (is_float)
                 {
-                    z_data[0] *= fbeta;
+                    z_data[0] = dbeta != 0 ? dbeta * z_data[0] : 0.f;
                     z_data[0] += alpha * sdot_(&NA1,  (float*)(A_data), &SA1,
                                               (float*)x_data, &Sx);
                 }
                 else
                 {
-                    z_data[0] *= dbeta;
+                    z_data[0] = dbeta != 0 ? dbeta * z_data[0] : 0.;
                     z_data[0] += alpha * ddot_(&NA1,  (double*)(A_data), &SA1,
                                               (double*)x_data, &Sx);
                 }
@@ -583,21 +592,21 @@ class CGemv(BaseBLAS, Gemv):
             alpha,
             beta,
             fail=sub["fail"],
-            force_init_beta=check_force_gemv_init(),
+            must_initialize_y=must_initialize_y_gemv(),
             params=sub["params"],
         )
         return code
 
     def c_code_cache_version(self):
-        return (16, blas_header_version(), check_force_gemv_init())
+        return (17, blas_header_version(), must_initialize_y_gemv())
 
 
 cgemv_inplace = CGemv(inplace=True)
 cgemv_no_inplace = CGemv(inplace=False)
 
 
-def check_force_gemv_init():
-    if check_force_gemv_init._force_init_beta is None:
+def must_initialize_y_gemv():
+    if must_initialize_y_gemv._force_init_beta is None:
         from pytensor.link.c.cmodule import GCC_compiler
 
         """
@@ -643,13 +652,13 @@ int main() {
         )
         if res:
             if res[0]:
-                check_force_gemv_init._force_init_beta = res[1]
+                must_initialize_y_gemv._force_init_beta = res[1]
             else:
-                check_force_gemv_init._force_init_beta = False
+                must_initialize_y_gemv._force_init_beta = False
         else:
-            check_force_gemv_init._force_init_beta = False
+            must_initialize_y_gemv._force_init_beta = False
 
-    return check_force_gemv_init._force_init_beta
+    return must_initialize_y_gemv._force_init_beta
 
 
-check_force_gemv_init._force_init_beta = None
+must_initialize_y_gemv._force_init_beta = None
diff --git a/pytensor/tensor/blockwise.py b/pytensor/tensor/blockwise.py
index 1c2a06992..29eafef99 100644
--- a/pytensor/tensor/blockwise.py
+++ b/pytensor/tensor/blockwise.py
@@ -16,6 +16,7 @@ from pytensor.graph.replace import (
     _vectorize_not_needed,
     vectorize_graph,
 )
+from pytensor.link.c.op import COp
 from pytensor.scalar import ScalarType
 from pytensor.tensor import as_tensor_variable
 from pytensor.tensor.shape import shape_padleft
@@ -43,7 +44,18 @@ def _vectorize_node_perform(
     """
 
     storage_map = {var: [None] for var in core_node.inputs + core_node.outputs}
-    core_thunk = core_node.op.make_thunk(core_node, storage_map, None, [], impl=impl)
+    try:
+        core_thunk = core_node.op.make_thunk(
+            core_node, storage_map, None, [], impl=impl
+        )
+    except NotImplementedError:
+        if impl == "c":
+            # Try again with py impl
+            core_thunk = core_node.op.make_thunk(
+                core_node, storage_map, None, [], impl="py"
+            )
+        else:
+            raise
     single_in = len(core_node.inputs) == 1
     core_input_storage = [storage_map[inp] for inp in core_node.inputs]
     core_output_storage = [storage_map[out] for out in core_node.outputs]
@@ -128,7 +140,7 @@ def _check_runtime_broadcast_core(numerical_inputs, batch_bcast_patterns, batch_
             )
 
 
-class Blockwise(Op):
+class Blockwise(COp):
     """Generalizes a core `Op` to work with batched dimensions.
 
     TODO: Dispatch JAX (should be easy with the vectorize macro)
@@ -483,6 +495,14 @@ class Blockwise(Op):
         else:
             return self.name
 
+    def c_code(self, *args, **kwargs):
+        # Blockwise is a C_Op just so we can propagate compilation mode to the inner Op.
+        # It doesn't itself have a C implementation yet.
+        raise NotImplementedError()
+
+    def c_code_cache_version(self):
+        return (-1,)
+
 
 @_vectorize_node.register(Op)
 def vectorize_node_fallback(op: Op, node: Apply, *bached_inputs) -> Apply:
diff --git a/pytensor/tensor/rewriting/blas.py b/pytensor/tensor/rewriting/blas.py
index b5c256448..e626b0720 100644
--- a/pytensor/tensor/rewriting/blas.py
+++ b/pytensor/tensor/rewriting/blas.py
@@ -700,7 +700,7 @@ def local_dot22_to_ger_or_gemv(fgraph, node):
             new_out = [rval]
         elif xb[0] and yb[1]:
             # x and y are both vectors so this qualifies for a sdot / ddot
-            # TODO: PyTensor doesn't have a sdot, but gemv is better than _dot22
+            # PyTensor's CGemv will call sdot/ddot at runtime, the Scipy Gemv may not
             xv = x.dimshuffle(1)
             zeros = ptb.AllocEmpty(x.dtype)(1)
             rval = gemv_no_inplace(zeros, one, y.T, xv, zero)
diff --git a/pytensor/tensor/signal/conv.py b/pytensor/tensor/signal/conv.py
index 59a5a9ea9..26c210fda 100644
--- a/pytensor/tensor/signal/conv.py
+++ b/pytensor/tensor/signal/conv.py
@@ -2,7 +2,8 @@ from typing import TYPE_CHECKING, Literal, cast
 
 from numpy import convolve as numpy_convolve
 
-from pytensor.graph import Apply, Op
+from pytensor.graph import Apply
+from pytensor.link.c.op import COp
 from pytensor.scalar.basic import upcast
 from pytensor.tensor.basic import as_tensor_variable, join, zeros
 from pytensor.tensor.blockwise import Blockwise
@@ -15,7 +16,7 @@ if TYPE_CHECKING:
     from pytensor.tensor import TensorLike
 
 
-class Convolve1d(Op):
+class Convolve1d(COp):
     __props__ = ("mode",)
     gufunc_signature = "(n),(k)->(o)"
 
@@ -86,6 +87,87 @@ class Convolve1d(Op):
 
         return [in1_bar, in2_bar]
 
+    def c_code_cache_version(self):
+        return (1,)
+
+    def c_code(self, node, name, inputs, outputs, sub):
+        # raise NotImplementedError()
+        in1, in2 = inputs
+        [out] = outputs
+        mode_str = self.mode
+
+        if mode_str == "full":
+            np_mode_val = 2  # NPY_CONVOLVE_FULL
+        elif mode_str == "valid":
+            np_mode_val = 0  # NPY_CONVOLVE_VALID
+        else:
+            # This case should ideally be prevented by __init__ or make_node
+            raise ValueError(f"Unsupported mode {mode_str}")
+
+        code = f"""
+        {{
+            PyArrayObject* in2_flipped_view = NULL;
+
+            if (PyArray_NDIM({in1}) != 1 || PyArray_NDIM({in2}) != 1) {{
+                PyErr_SetString(PyExc_ValueError, "Convolve1d C code expects 1D arrays.");
+                {sub['fail']};
+            }}
+
+            npy_intp n_in2 = PyArray_DIM({in2}, 0);
+
+            // Create a reversed view of in2
+            if (n_in2 == 0) {{
+                PyErr_SetString(PyExc_ValueError, "Convolve1d: second input (kernel) cannot be empty.");
+                {sub['fail']};
+            }} else {{
+                npy_intp view_dims[1];
+                view_dims[0] = n_in2;
+
+                npy_intp view_strides[1];
+                view_strides[0] = -PyArray_STRIDES({in2})[0];
+
+                void* view_data = (char*)PyArray_DATA({in2}) + (n_in2 - 1) * PyArray_STRIDES({in2})[0];
+
+                Py_INCREF(PyArray_DESCR({in2}));
+                in2_flipped_view = (PyArrayObject*)PyArray_NewFromDescr(
+                    Py_TYPE({in2}),
+                    PyArray_DESCR({in2}),
+                    1,  // ndim
+                    view_dims,
+                    view_strides,
+                    view_data,
+                    (PyArray_FLAGS({in2}) & ~NPY_ARRAY_WRITEABLE),
+                    NULL
+                );
+
+                if (!in2_flipped_view) {{
+                    PyErr_SetString(PyExc_RuntimeError, "Failed to create flipped kernel view for Convolve1d.");
+                    {sub['fail']};
+                }}
+
+                Py_INCREF({in2});
+                if (PyArray_SetBaseObject(in2_flipped_view, (PyObject*){in2}) < 0) {{
+                    Py_DECREF({in2}); // SetBaseObject failed, release the extra INCREF
+                    Py_DECREF(in2_flipped_view);
+                    in2_flipped_view = NULL;
+                    PyErr_SetString(PyExc_RuntimeError, "Failed to set base object for flipped kernel view in Convolve1d.");
+                    {sub['fail']};
+                }}
+                PyArray_UpdateFlags(in2_flipped_view, (NPY_ARRAY_C_CONTIGUOUS | NPY_ARRAY_F_CONTIGUOUS));
+            }}
+
+            // TODO: Use lower level implementation that allows reusing the output buffer
+            Py_XDECREF({out});
+            {out} = (PyArrayObject*) PyArray_Correlate2((PyObject*){in1}, (PyObject*)in2_flipped_view, {np_mode_val});
+            Py_XDECREF(in2_flipped_view); // Clean up the view if correlate fails
+            if (!{out}) {{
+                // PyArray_Correlate already set an error
+                {sub['fail']};
+            }}
+        }}
+        """
+        return code
+
 
 def convolve1d(
     in1: "TensorLike",
diff --git a/pytensor/tensor/type.py b/pytensor/tensor/type.py
index 0474aad77..5ae92006e 100644
--- a/pytensor/tensor/type.py
+++ b/pytensor/tensor/type.py
@@ -793,7 +793,7 @@ def tensor(
         try:
             # Help catching errors with the new tensor API
             # Many single letter strings are valid sctypes
-            if str(name) == "floatX" or (len(str(name)) > 1 and np.dtype(name).type):
+            if str(name) == "floatX" or (len(str(name)) > 2 and np.dtype(name).type):
                 raise ValueError(
                     f"The first and only positional argument of tensor is now `name`. Got {name}.\n"
                     "This name looks like a dtype, which you should pass as a keyword argument only."


Test patch:
diff --git a/tests/link/jax/test_slinalg.py b/tests/link/jax/test_slinalg.py
index ca944221a..b2b722f8b 100644
--- a/tests/link/jax/test_slinalg.py
+++ b/tests/link/jax/test_slinalg.py
@@ -122,6 +122,54 @@ def test_jax_solve():
     )
 
 
+@pytest.mark.parametrize(
+    "A_size, b_size, b_ndim",
+    [
+        (
+            (5, 5),
+            (5,),
+            1,
+        ),
+        (
+            (5, 5),
+            (5, 1),
+            2,
+        ),
+        (
+            (5, 5),
+            (1, 5),
+            1,
+        ),
+        (
+            (4, 5, 5),
+            (4, 5, 5),
+            2,
+        ),
+    ],
+    ids=["basic_vector", "basic_matrix", "vector_broadcasted", "fully_batched"],
+)
+def test_jax_tridiagonal_solve(A_size: tuple, b_size: tuple, b_ndim: int):
+    A = pt.tensor("A", shape=A_size)
+    b = pt.tensor("b", shape=b_size)
+
+    out = pt.linalg.solve(A, b, assume_a="tridiagonal", b_ndim=b_ndim)
+
+    A_val = np.zeros(A_size)
+    N = A_size[-1]
+    A_val[...] = np.eye(N)
+    for i in range(N - 1):
+        A_val[..., i, i + 1] = np.random.randn()
+        A_val[..., i + 1, i] = np.random.randn()
+
+    b_val = np.random.randn(*b_size)
+
+    compare_jax_and_py(
+        [A, b],
+        [out],
+        [A_val, b_val],
+    )
+
+
 def test_jax_SolveTriangular():
     rng = np.random.default_rng(utt.fetch_seed())
 
diff --git a/tests/link/numba/test_basic.py b/tests/link/numba/test_basic.py
index e8390b8eb..9132d7b20 100644
--- a/tests/link/numba/test_basic.py
+++ b/tests/link/numba/test_basic.py
@@ -14,7 +14,6 @@ from tests.tensor.test_math_scipy import scipy
 numba = pytest.importorskip("numba")
 
 import pytensor.scalar as ps
-import pytensor.scalar.math as psm
 import pytensor.tensor as pt
 import pytensor.tensor.math as ptm
 from pytensor import config, shared
@@ -260,9 +259,12 @@ def compare_numba_and_py(
     if assert_fn is None:
 
         def assert_fn(x, y):
-            return np.testing.assert_allclose(x, y, rtol=1e-4) and compare_shape_dtype(
-                x, y
-            )
+            np.testing.assert_allclose(x, y, rtol=1e-4, strict=True)
+            # Make sure we don't have one input be a np.ndarray while the other is not
+            if isinstance(x, np.ndarray):
+                assert isinstance(y, np.ndarray), "y is not a NumPy array, but x is"
+            else:
+                assert not isinstance(y, np.ndarray), "y is a NumPy array, but x is not"
 
     if any(
         inp.owner is not None
@@ -295,8 +297,8 @@ 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)
+        for numba_res_i, python_res_i in zip(numba_res, py_res, strict=True):
+            assert_fn(numba_res_i, python_res_i)
     else:
         assert_fn(numba_res, py_res)
 
@@ -640,48 +642,6 @@ def test_Dot(x, y, exc):
         )
 
 
-@pytest.mark.parametrize(
-    "x, exc",
-    [
-        (
-            (ps.float64(), np.array(0.0, dtype="float64")),
-            None,
-        ),
-        (
-            (ps.float64(), np.array(-32.0, dtype="float64")),
-            None,
-        ),
-        (
-            (ps.float64(), np.array(-40.0, dtype="float64")),
-            None,
-        ),
-        (
-            (ps.float64(), np.array(32.0, dtype="float64")),
-            None,
-        ),
-        (
-            (ps.float64(), np.array(40.0, dtype="float64")),
-            None,
-        ),
-        (
-            (ps.int64(), np.array(32, dtype="int64")),
-            None,
-        ),
-    ],
-)
-def test_Softplus(x, exc):
-    x, x_test_value = x
-    g = psm.Softplus(ps.upgrade_to_float)(x)
-
-    cm = contextlib.suppress() if exc is None else pytest.warns(exc)
-    with cm:
-        compare_numba_and_py(
-            [x],
-            [g],
-            [x_test_value],
-        )
-
-
 @pytest.mark.parametrize(
     "x, y, exc",
     [
diff --git a/tests/link/numba/test_nlinalg.py b/tests/link/numba/test_nlinalg.py
index 67bdc6f1a..8d7c3a449 100644
--- a/tests/link/numba/test_nlinalg.py
+++ b/tests/link/numba/test_nlinalg.py
@@ -11,68 +11,20 @@ from tests.link.numba.test_basic import compare_numba_and_py
 rng = np.random.default_rng(42849)
 
 
+@pytest.mark.parametrize("dtype", ("float64", "int64"))
 @pytest.mark.parametrize(
-    "x, exc",
-    [
-        (
-            (
-                pt.dmatrix(),
-                (lambda x: x.T.dot(x))(rng.random(size=(3, 3)).astype("float64")),
-            ),
-            None,
-        ),
-        (
-            (
-                pt.lmatrix(),
-                (lambda x: x.T.dot(x))(rng.poisson(size=(3, 3)).astype("int64")),
-            ),
-            None,
-        ),
-    ],
+    "op", (nlinalg.Det(), nlinalg.SLogDet()), ids=["det", "slogdet"]
 )
-def test_Det(x, exc):
-    x, test_x = x
-    g = nlinalg.Det()(x)
+def test_Det_SLogDet(op, dtype):
+    x = pt.matrix(dtype=dtype)
 
-    cm = contextlib.suppress() if exc is None else pytest.warns(exc)
-    with cm:
-        compare_numba_and_py(
-            [x],
-            g,
-            [test_x],
-        )
+    rng = np.random.default_rng([50, sum(map(ord, dtype))])
+    x_ = rng.random(size=(3, 3)).astype(dtype)
+    test_x = x_.T.dot(x_)
 
+    g = op(x)
 
-@pytest.mark.parametrize(
-    "x, exc",
-    [
-        (
-            (
-                pt.dmatrix(),
-                (lambda x: x.T.dot(x))(rng.random(size=(3, 3)).astype("float64")),
-            ),
-            None,
-        ),
-        (
-            (
-                pt.lmatrix(),
-                (lambda x: x.T.dot(x))(rng.poisson(size=(3, 3)).astype("int64")),
-            ),
-            None,
-        ),
-    ],
-)
-def test_SLogDet(x, exc):
-    x, test_x = x
-    g = nlinalg.SLogDet()(x)
-
-    cm = contextlib.suppress() if exc is None else pytest.warns(exc)
-    with cm:
-        compare_numba_and_py(
-            [x],
-            g,
-            [test_x],
-        )
+    compare_numba_and_py([x], g, [test_x])
 
 
 # We were seeing some weird results in CI where the following two almost
diff --git a/tests/link/numba/test_scalar.py b/tests/link/numba/test_scalar.py
index 504d2a163..2125d7cc0 100644
--- a/tests/link/numba/test_scalar.py
+++ b/tests/link/numba/test_scalar.py
@@ -3,12 +3,13 @@ import pytest
 
 import pytensor.scalar as ps
 import pytensor.scalar.basic as psb
+import pytensor.scalar.math as psm
 import pytensor.tensor as pt
-from pytensor import config
+from pytensor import config, function
 from pytensor.scalar.basic import Composite
 from pytensor.tensor import tensor
 from pytensor.tensor.elemwise import Elemwise
-from tests.link.numba.test_basic import compare_numba_and_py
+from tests.link.numba.test_basic import compare_numba_and_py, numba_mode, py_mode
 
 
 rng = np.random.default_rng(42849)
@@ -99,7 +100,11 @@ def test_Composite(inputs, input_values, scalar_fn):
     "v, dtype",
     [
         ((pt.fscalar(), np.array(1.0, dtype="float32")), psb.float64),
-        ((pt.dscalar(), np.array(1.0, dtype="float64")), psb.float32),
+        pytest.param(
+            (pt.dscalar(), np.array(1.0, dtype="float64")),
+            psb.float32,
+            marks=pytest.mark.xfail(reason="Scalar downcasting not supported in numba"),
+        ),
     ],
 )
 def test_Cast(v, dtype):
@@ -145,3 +150,37 @@ def test_isnan(composite):
         [out],
         [np.array([1, 0], dtype="float64")],
     )
+
+
+@pytest.mark.parametrize(
+    "dtype",
+    [
+        pytest.param(
+            "float32",
+            marks=pytest.mark.xfail(reason="Scalar downcasting not supported in numba"),
+        ),
+        "float64",
+        pytest.param(
+            "int16",
+            marks=pytest.mark.xfail(reason="Scalar downcasting not supported in numba"),
+        ),
+        "int64",
+        "uint32",
+    ],
+)
+def test_Softplus(dtype):
+    x = ps.get_scalar_type(dtype)("x")
+    g = psm.softplus(x)
+
+    py_fn = function([x], g, mode=py_mode)
+    numba_fn = function([x], g, mode=numba_mode)
+    for value in (-40, -32, 0, 32, 40):
+        if value < 0 and dtype.startswith("u"):
+            continue
+        test_x = np.dtype(dtype).type(value)
+        np.testing.assert_allclose(
+            py_fn(test_x),
+            numba_fn(test_x),
+            strict=True,
+            err_msg=f"Failed for value {value}",
+        )
diff --git a/tests/tensor/linalg/test_rewriting.py b/tests/tensor/linalg/test_rewriting.py
index 6f04fac5f..32683029f 100644
--- a/tests/tensor/linalg/test_rewriting.py
+++ b/tests/tensor/linalg/test_rewriting.py
@@ -161,3 +161,36 @@ def test_lu_decomposition_reused_scan(transposed):
     resx1 = fn_opt(A_test, x0_test)
     rtol = 1e-7 if config.floatX == "float64" else 1e-6
     np.testing.assert_allclose(resx0, resx1, rtol=rtol)
+
+
+def test_lu_decomposition_reused_preserves_check_finite():
+    # Check that the LU decomposition rewrite preserves the check_finite flag
+    rewrite_name = reuse_lu_decomposition_multiple_solves.__name__
+
+    A = tensor("A", shape=(2, 2))
+    b1 = tensor("b1", shape=(2,))
+    b2 = tensor("b2", shape=(2,))
+
+    x1 = solve(A, b1, assume_a="gen", check_finite=True)
+    x2 = solve(A, b2, assume_a="gen", check_finite=False)
+    fn_opt = function(
+        [A, b1, b2], [x1, x2], mode=get_default_mode().including(rewrite_name)
+    )
+    opt_nodes = fn_opt.maker.fgraph.apply_nodes
+    assert count_vanilla_solve_nodes(opt_nodes) == 0
+    assert count_lu_decom_nodes(opt_nodes) == 1
+    assert count_lu_solve_nodes(opt_nodes) == 2
+
+    # We should get an error if A or b1 is non finite
+    A_valid = np.array([[1, 0], [0, 1]], dtype=A.type.dtype)
+    b1_valid = np.array([1, 1], dtype=b1.type.dtype)
+    b2_valid = np.array([1, 1], dtype=b2.type.dtype)
+
+    assert fn_opt(A_valid, b1_valid, b2_valid)  # Fine
+    assert fn_opt(
+        A_valid, b1_valid, b2_valid * np.nan
+    )  # Should not raise (also fine on most LAPACK implementations?)
+    with pytest.raises(ValueError, match="array must not contain infs or NaNs"):
+        assert fn_opt(A_valid, b1_valid * np.nan, b2_valid)
+    with pytest.raises(ValueError, match="array must not contain infs or NaNs"):
+        assert fn_opt(A_valid * np.nan, b1_valid, b2_valid)
diff --git a/tests/tensor/test_blas_c.py b/tests/tensor/test_blas_c.py
index e46c03676..b6ba1987b 100644
--- a/tests/tensor/test_blas_c.py
+++ b/tests/tensor/test_blas_c.py
@@ -7,7 +7,7 @@ import pytensor
 import pytensor.tensor as pt
 from pytensor.tensor.basic import AllocEmpty
 from pytensor.tensor.blas import Ger
-from pytensor.tensor.blas_c import CGemv, CGer, check_force_gemv_init
+from pytensor.tensor.blas_c import CGemv, CGer, must_initialize_y_gemv
 from pytensor.tensor.blas_scipy import ScipyGer
 from pytensor.tensor.type import dmatrix, dvector, matrix, scalar, tensor, vector
 from tests import unittest_tools
@@ -130,31 +130,35 @@ class TestCGemv(OptimizationTestMixin):
         self.dtype = dtype
         self.mode = pytensor.compile.get_default_mode().including("fast_run")
         # matrix
-        self.A = tensor(dtype=dtype, shape=(None, None))
+        self.A = tensor("A", dtype=dtype, shape=(None, None))
         self.Aval = np.ones((2, 3), dtype=dtype)
 
         # vector
-        self.x = tensor(dtype=dtype, shape=(None,))
-        self.y = tensor(dtype=dtype, shape=(None,))
+        self.x = tensor("x", dtype=dtype, shape=(None,))
+        self.y = tensor("y", dtype=dtype, shape=(None,))
         self.xval = np.asarray([1, 2], dtype=dtype)
         self.yval = np.asarray([1.5, 2.7, 3.9], dtype=dtype)
 
         # scalar
-        self.a = tensor(dtype=dtype, shape=())
+        self.a = tensor("a", dtype=dtype, shape=())
 
-    def test_nan_beta_0(self):
+    @pytest.mark.parametrize("inplace", [True, False])
+    def test_nan_beta_0(self, inplace):
         mode = self.mode.including()
         mode.check_isfinite = False
         f = pytensor.function(
-            [self.A, self.x, self.y, self.a],
+            [self.A, self.x, pytensor.In(self.y, mutable=inplace), self.a],
             self.a * self.y + pt.dot(self.A, self.x),
             mode=mode,
         )
-        Aval = np.ones((3, 1), dtype=self.dtype)
-        xval = np.ones((1,), dtype=self.dtype)
-        yval = float("NaN") * np.ones((3,), dtype=self.dtype)
-        zval = f(Aval, xval, yval, 0)
-        assert not np.isnan(zval).any()
+        [node] = f.maker.fgraph.apply_nodes
+        assert isinstance(node.op, CGemv) and node.op.inplace == inplace
+        for rows in (3, 1):
+            Aval = np.ones((rows, 1), dtype=self.dtype)
+            xval = np.ones((1,), dtype=self.dtype)
+            yval = np.full((rows,), np.nan, dtype=self.dtype)
+            zval = f(Aval, xval, yval, 0)
+            assert not np.isnan(zval).any()
 
     def test_optimizations_vm(self):
         skip_if_blas_ldflags_empty()
@@ -191,8 +195,10 @@ class TestCGemv(OptimizationTestMixin):
             np.dot(self.Aval[::-1, ::-1], self.yval),
         )
 
-    def test_force_gemv_init(self):
-        if check_force_gemv_init():
+    def test_must_initialize_y_gemv(self):
+        if must_initialize_y_gemv():
+            # FIME: This warn should be emitted by the function if we find it relevant
+            # Not in a test that doesn't care about the outcome either way
             warn(
                 "WARNING: The current BLAS requires PyTensor to initialize"
                 " memory for some GEMV calls which will result in a minor"