Repository: pymc-devs/pytensor

Issue:
Einsum fails with repeated inputs
### Description

```python
import pytensor.tensor as pt

x = pt.matrix("x")
pt.einsum("ij,ij->i", x, x.copy())  # Fine
pt.einsum("ij,ij->i", x, x)  # Raises
```
```
ValueError: The following variables were provided more than once as inputs to the OpFromGraph, resulting in an invalid graph: [x]. Use dummy variables or var.copy() to distinguish variables when creating the OpFromGraph graph.
```
We should add an `x.copy()` internally, since this is an implementation detail

Discussion:

Solution patch:
diff --git a/.github/workflows/pypi.yml b/.github/workflows/pypi.yml
index e588a5eae..3e5f674a5 100644
--- a/.github/workflows/pypi.yml
+++ b/.github/workflows/pypi.yml
@@ -80,7 +80,7 @@ jobs:
           persist-credentials: false
 
       - name: Build wheels
-        uses: pypa/cibuildwheel@v2.22.0
+        uses: pypa/cibuildwheel@v2.23.0
 
       - name: Attest GitHub build provenance
         uses: actions/attest-build-provenance@v2
diff --git a/pytensor/link/jax/dispatch/slinalg.py b/pytensor/link/jax/dispatch/slinalg.py
index ca362e453..5430ce1da 100644
--- a/pytensor/link/jax/dispatch/slinalg.py
+++ b/pytensor/link/jax/dispatch/slinalg.py
@@ -53,7 +53,6 @@ def jax_funcify_Solve(op, **kwargs):
 @jax_funcify.register(SolveTriangular)
 def jax_funcify_SolveTriangular(op, **kwargs):
     lower = op.lower
-    trans = op.trans
     unit_diagonal = op.unit_diagonal
     check_finite = op.check_finite
 
@@ -62,7 +61,7 @@ def jax_funcify_SolveTriangular(op, **kwargs):
             A,
             b,
             lower=lower,
-            trans=trans,
+            trans=0,  # this is handled by explicitly transposing A, so it will always be 0 when we get to here.
             unit_diagonal=unit_diagonal,
             check_finite=check_finite,
         )
diff --git a/pytensor/link/numba/dispatch/basic.py b/pytensor/link/numba/dispatch/basic.py
index 7685c17d9..19e91e5f8 100644
--- a/pytensor/link/numba/dispatch/basic.py
+++ b/pytensor/link/numba/dispatch/basic.py
@@ -243,93 +243,6 @@ def impl_to_scalar(x):
         raise TypingError(f"{x} must be a scalar compatible type.")
 
 
-def enable_slice_literals():
-    """Enable lowering for ``SliceLiteral``s.
-
-    TODO: This can be removed once https://github.com/numba/numba/pull/6996 is merged
-    and a release is made.
-    """
-    from numba.core import types
-    from numba.core.datamodel.models import SliceModel
-    from numba.core.datamodel.registry import register_default
-    from numba.core.imputils import lower_cast, lower_constant
-    from numba.core.types.misc import SliceLiteral
-    from numba.cpython.slicing import get_defaults
-
-    register_default(numba.types.misc.SliceLiteral)(SliceModel)
-
-    @property
-    def key(self):
-        return self.name
-
-    SliceLiteral.key = key
-
-    def make_slice_from_constant(context, builder, ty, pyval):
-        sli = context.make_helper(builder, ty)
-        lty = context.get_value_type(types.intp)
-
-        (
-            default_start_pos,
-            default_start_neg,
-            default_stop_pos,
-            default_stop_neg,
-            default_step,
-        ) = (context.get_constant(types.intp, x) for x in get_defaults(context))
-
-        step = pyval.step
-        if step is None:
-            step_is_neg = False
-            step = default_step
-        else:
-            step_is_neg = step < 0
-            step = lty(step)
-
-        start = pyval.start
-        if start is None:
-            if step_is_neg:
-                start = default_start_neg
-            else:
-                start = default_start_pos
-        else:
-            start = lty(start)
-
-        stop = pyval.stop
-        if stop is None:
-            if step_is_neg:
-                stop = default_stop_neg
-            else:
-                stop = default_stop_pos
-        else:
-            stop = lty(stop)
-
-        sli.start = start
-        sli.stop = stop
-        sli.step = step
-
-        return sli._getvalue()
-
-    @lower_constant(numba.types.SliceType)
-    def constant_slice(context, builder, ty, pyval):
-        if isinstance(ty, types.Literal):
-            typ = ty.literal_type
-        else:
-            typ = ty
-
-        return make_slice_from_constant(context, builder, typ, pyval)
-
-    @lower_cast(numba.types.misc.SliceLiteral, numba.types.SliceType)
-    def cast_from_literal(context, builder, fromty, toty, val):
-        return make_slice_from_constant(
-            context,
-            builder,
-            toty,
-            fromty.literal_value,
-        )
-
-
-enable_slice_literals()
-
-
 def create_tuple_creator(f, n):
     """Construct a compile-time ``tuple``-comprehension-like loop.
 
diff --git a/pytensor/link/numba/dispatch/slinalg.py b/pytensor/link/numba/dispatch/slinalg.py
index 4b5f51892..c64b5fdb3 100644
--- a/pytensor/link/numba/dispatch/slinalg.py
+++ b/pytensor/link/numba/dispatch/slinalg.py
@@ -180,7 +180,6 @@ def solve_triangular_impl(A, B, trans, lower, unit_diagonal, b_ndim, overwrite_b
 
 @numba_funcify.register(SolveTriangular)
 def numba_funcify_SolveTriangular(op, node, **kwargs):
-    trans = bool(op.trans)
     lower = op.lower
     unit_diagonal = op.unit_diagonal
     check_finite = op.check_finite
@@ -208,7 +207,7 @@ def numba_funcify_SolveTriangular(op, node, **kwargs):
         res = _solve_triangular(
             a,
             b,
-            trans=trans,
+            trans=0,  # transposing is handled explicitly on the graph, so we never use this argument
             lower=lower,
             unit_diagonal=unit_diagonal,
             overwrite_b=overwrite_b,
diff --git a/pytensor/link/numba/dispatch/subtensor.py b/pytensor/link/numba/dispatch/subtensor.py
index 94cb51434..81348b57b 100644
--- a/pytensor/link/numba/dispatch/subtensor.py
+++ b/pytensor/link/numba/dispatch/subtensor.py
@@ -150,7 +150,7 @@ def numba_funcify_AdvancedSubtensor(op, node, **kwargs):
             for adv_idx in adv_idxs
         )
         # Must be consecutive
-        and not op.non_contiguous_adv_indexing(node)
+        and not op.non_consecutive_adv_indexing(node)
         # y in set/inc_subtensor cannot be broadcasted
         and (
             y is None
diff --git a/pytensor/tensor/c_code/dimshuffle.c b/pytensor/tensor/c_code/dimshuffle.c
index b99a0ee41..0bfc5df3b 100644
--- a/pytensor/tensor/c_code/dimshuffle.c
+++ b/pytensor/tensor/c_code/dimshuffle.c
@@ -33,12 +33,17 @@ int APPLY_SPECIFIC(cpu_dimshuffle)(PyArrayObject *input, PyArrayObject **res, PA
     npy_intp original_size = PyArray_SIZE(input);
     npy_intp new_size = 1;
     for (npy_intp i = 0; i < nd_out; ++i) {
+        // We set the strides of length 1 dimensions to PyArray_ITEMSIZE(input).
+        // The value is arbitrary, because there is never a next element.
+        // np.expand_dims(x, 0) and x[None] do different things here.
+        // I would prefer zero, but there are some poorly implemented BLAS operations
+        // That don't handle zero strides correctly. At least they won't fail because of DimShuffle.
         if (new_order[i] != -1) {
             dimensions[i] = PyArray_DIMS(input)[new_order[i]];
-            strides[i] = PyArray_DIMS(input)[new_order[i]] == 1 ? 0 : PyArray_STRIDES(input)[new_order[i]];
+            strides[i] = PyArray_DIMS(input)[new_order[i]] == 1 ? PyArray_ITEMSIZE(input) : PyArray_STRIDES(input)[new_order[i]];
         } else {
             dimensions[i] = 1;
-            strides[i] = 0;
+            strides[i] = PyArray_ITEMSIZE(input);
         }
         new_size *= dimensions[i];
     }
diff --git a/pytensor/tensor/einsum.py b/pytensor/tensor/einsum.py
index 660c16d38..e119b6de1 100644
--- a/pytensor/tensor/einsum.py
+++ b/pytensor/tensor/einsum.py
@@ -417,6 +417,18 @@ def _right_to_left_path(n: int) -> tuple[tuple[int, int], ...]:
     return tuple(pairwise(reversed(range(n))))
 
 
+def _ensure_not_equal(elements):
+    """
+    Ensures that any pair in a list of elements are not the same object. If a pair of elements is found to be equal, then one of them is converted to a copy.
+    """
+    elements = list(elements)
+    for i, elem1 in enumerate(elements[:-1]):
+        for j, elem2 in enumerate(elements[i + 1 :], start=i + 1):
+            if elem1 is elem2:
+                elements[j] = elem1.copy()
+    return elements
+
+
 def einsum(subscripts: str, *operands: "TensorLike", optimize=None) -> TensorVariable:
     """
     Multiplication and summation of tensors using the Einstein summation convention.
@@ -553,7 +565,7 @@ def einsum(subscripts: str, *operands: "TensorLike", optimize=None) -> TensorVar
             "If you need this functionality open an issue in https://github.com/pymc-devs/pytensor/issues to let us know. "
         )
 
-    tensor_operands = [as_tensor(operand) for operand in operands]
+    tensor_operands = _ensure_not_equal([as_tensor(operand) for operand in operands])
     shapes = [operand.type.shape for operand in tensor_operands]
 
     path: PATH
diff --git a/pytensor/tensor/math.py b/pytensor/tensor/math.py
index af6a3827a..e4d73aa63 100644
--- a/pytensor/tensor/math.py
+++ b/pytensor/tensor/math.py
@@ -4122,6 +4122,154 @@ def matmul(x1: "ArrayLike", x2: "ArrayLike", dtype: Optional["DTypeLike"] = None
     return out
 
 
+def vecdot(
+    x1: TensorLike,
+    x2: TensorLike,
+    dtype: Optional["DTypeLike"] = None,
+) -> TensorVariable:
+    """Compute the vector dot product of two arrays.
+
+    Parameters
+    ----------
+    x1, x2
+        Input arrays with the same shape.
+    dtype
+        The desired data-type for the result. If not given, then the type will
+        be determined as the minimum type required to hold the objects in the
+        sequence.
+
+    Returns
+    -------
+    TensorVariable
+        The vector dot product of the inputs.
+
+    Notes
+    -----
+    This is equivalent to `numpy.vecdot` and computes the dot product of
+    vectors along the last axis of both inputs. Broadcasting is supported
+    across all other dimensions.
+
+    Examples
+    --------
+    >>> import pytensor.tensor as pt
+    >>> # Vector dot product with shape (5,) inputs
+    >>> x = pt.vector("x", shape=(5,))  # shape (5,)
+    >>> y = pt.vector("y", shape=(5,))  # shape (5,)
+    >>> z = pt.vecdot(x, y)  # scalar output
+    >>> # Equivalent to numpy.vecdot(x, y)
+    >>>
+    >>> # With batched inputs of shape (3, 5)
+    >>> x_batch = pt.matrix("x", shape=(3, 5))  # shape (3, 5)
+    >>> y_batch = pt.matrix("y", shape=(3, 5))  # shape (3, 5)
+    >>> z_batch = pt.vecdot(x_batch, y_batch)  # shape (3,)
+    >>> # Equivalent to numpy.vecdot(x_batch, y_batch)
+    """
+    out = _inner_prod(x1, x2)
+
+    if dtype is not None:
+        out = out.astype(dtype)
+
+    return out
+
+
+def matvec(
+    x1: TensorLike, x2: TensorLike, dtype: Optional["DTypeLike"] = None
+) -> TensorVariable:
+    """Compute the matrix-vector product.
+
+    Parameters
+    ----------
+    x1
+        Input array for the matrix with shape (..., M, K).
+    x2
+        Input array for the vector with shape (..., K).
+    dtype
+        The desired data-type for the result. If not given, then the type will
+        be determined as the minimum type required to hold the objects in the
+        sequence.
+
+    Returns
+    -------
+    TensorVariable
+        The matrix-vector product with shape (..., M).
+
+    Notes
+    -----
+    This is equivalent to `numpy.matvec` and computes the matrix-vector product
+    with broadcasting over batch dimensions.
+
+    Examples
+    --------
+    >>> import pytensor.tensor as pt
+    >>> # Matrix-vector product
+    >>> A = pt.matrix("A", shape=(3, 4))  # shape (3, 4)
+    >>> v = pt.vector("v", shape=(4,))  # shape (4,)
+    >>> result = pt.matvec(A, v)  # shape (3,)
+    >>> # Equivalent to numpy.matvec(A, v)
+    >>>
+    >>> # Batched matrix-vector product
+    >>> batched_A = pt.tensor3("A", shape=(2, 3, 4))  # shape (2, 3, 4)
+    >>> batched_v = pt.matrix("v", shape=(2, 4))  # shape (2, 4)
+    >>> result = pt.matvec(batched_A, batched_v)  # shape (2, 3)
+    >>> # Equivalent to numpy.matvec(batched_A, batched_v)
+    """
+    out = _matrix_vec_prod(x1, x2)
+
+    if dtype is not None:
+        out = out.astype(dtype)
+
+    return out
+
+
+def vecmat(
+    x1: TensorLike, x2: TensorLike, dtype: Optional["DTypeLike"] = None
+) -> TensorVariable:
+    """Compute the vector-matrix product.
+
+    Parameters
+    ----------
+    x1
+        Input array for the vector with shape (..., K).
+    x2
+        Input array for the matrix with shape (..., K, N).
+    dtype
+        The desired data-type for the result. If not given, then the type will
+        be determined as the minimum type required to hold the objects in the
+        sequence.
+
+    Returns
+    -------
+    TensorVariable
+        The vector-matrix product with shape (..., N).
+
+    Notes
+    -----
+    This is equivalent to `numpy.vecmat` and computes the vector-matrix product
+    with broadcasting over batch dimensions.
+
+    Examples
+    --------
+    >>> import pytensor.tensor as pt
+    >>> # Vector-matrix product
+    >>> v = pt.vector("v", shape=(3,))  # shape (3,)
+    >>> A = pt.matrix("A", shape=(3, 4))  # shape (3, 4)
+    >>> result = pt.vecmat(v, A)  # shape (4,)
+    >>> # Equivalent to numpy.vecmat(v, A)
+    >>>
+    >>> # Batched vector-matrix product
+    >>> batched_v = pt.matrix("v", shape=(2, 3))  # shape (2, 3)
+    >>> batched_A = pt.tensor3("A", shape=(2, 3, 4))  # shape (2, 3, 4)
+    >>> result = pt.vecmat(batched_v, batched_A)  # shape (2, 4)
+    >>> # Equivalent to numpy.vecmat(batched_v, batched_A)
+    """
+    out = _vec_matrix_prod(x1, x2)
+
+    if dtype is not None:
+        out = out.astype(dtype)
+
+    return out
+
+
 @_vectorize_node.register(Dot)
 def vectorize_node_dot(op, node, batched_x, batched_y):
     old_x, old_y = node.inputs
@@ -4218,6 +4366,9 @@ __all__ = [
     "max_and_argmax",
     "max",
     "matmul",
+    "vecdot",
+    "matvec",
+    "vecmat",
     "argmax",
     "min",
     "argmin",
diff --git a/pytensor/tensor/rewriting/subtensor.py b/pytensor/tensor/rewriting/subtensor.py
index 020a2e04e..c7a4574a9 100644
--- a/pytensor/tensor/rewriting/subtensor.py
+++ b/pytensor/tensor/rewriting/subtensor.py
@@ -2029,18 +2029,41 @@ def ravel_multidimensional_bool_idx(fgraph, node):
     return [copy_stack_trace(node.outputs[0], new_out)]
 
 
-@node_rewriter(tracks=[AdvancedSubtensor])
+@node_rewriter(tracks=[AdvancedSubtensor, AdvancedIncSubtensor])
 def ravel_multidimensional_int_idx(fgraph, node):
-    """Convert multidimensional integer indexing into equivalent vector integer index, supported by Numba
-
-    x[eye(3, dtype=int)] -> x[eye(3).ravel()].reshape((3, 3))
+    """Convert multidimensional integer indexing into equivalent consecutive vector integer index,
+    supported by Numba or by our specialized dispatchers
 
+        x[eye(3)] -> x[eye(3).ravel()].reshape((3, 3))
 
     NOTE: This is very similar to the rewrite `local_replace_AdvancedSubtensor` except it also handles non-full slices
 
-    x[eye(3, dtype=int), 2:] -> x[eye(3).ravel(), 2:].reshape((3, 3, ...)), where ... are the remaining output shapes
+        x[eye(3), 2:] -> x[eye(3).ravel(), 2:].reshape((3, 3, ...)), where ... are the remaining output shapes
+
+    It also handles multiple integer indices, but only if they don't broadcast
+
+        x[eye(3,), 2:, eye(3)] -> x[eye(3).ravel(), eye(3).ravel(), 2:].reshape((3, 3, ...)), where ... are the remaining output shapes
+
+    Also handles AdvancedIncSubtensor, but only if the advanced indices are consecutive and neither indices nor y broadcast
+
+        x[eye(3), 2:].set(y) -> x[eye(3).ravel(), 2:].set(y.reshape(-1, y.shape[1:]))
+
     """
-    x, *idxs = node.inputs
+    op = node.op
+    non_consecutive_adv_indexing = op.non_consecutive_adv_indexing(node)
+    is_inc_subtensor = isinstance(op, AdvancedIncSubtensor)
+
+    if is_inc_subtensor:
+        x, y, *idxs = node.inputs
+        # Inc/SetSubtensor is harder to reason about due to y
+        # We get out if it's broadcasting or if the advanced indices are non-consecutive
+        if non_consecutive_adv_indexing or (
+            y.type.broadcastable != x[tuple(idxs)].type.broadcastable
+        ):
+            return None
+
+    else:
+        x, *idxs = node.inputs
 
     if any(
         (
@@ -2049,39 +2072,90 @@ def ravel_multidimensional_int_idx(fgraph, node):
         )
         for idx in idxs
     ):
-        # Get out if there are any other advanced indexes or np.newaxis
+        # Get out if there are any other advanced indices or np.newaxis
         return None
 
-    int_idxs = [
+    int_idxs_and_pos = [
         (i, idx)
         for i, idx in enumerate(idxs)
         if (isinstance(idx.type, TensorType) and idx.dtype in integer_dtypes)
     ]
 
-    if len(int_idxs) != 1:
-        # Get out if there are no or multiple integer idxs
+    if not int_idxs_and_pos:
         return None
 
-    [(int_idx_pos, int_idx)] = int_idxs
-    if int_idx.type.ndim < 2:
-        # No need to do anything if it's a vector or scalar, as it's already supported by Numba
+    int_idxs_pos, int_idxs = zip(
+        *int_idxs_and_pos, strict=False
+    )  # strict=False because by definition it's true
+
+    first_int_idx_pos = int_idxs_pos[0]
+    first_int_idx = int_idxs[0]
+    first_int_idx_bcast = first_int_idx.type.broadcastable
+
+    if any(int_idx.type.broadcastable != first_int_idx_bcast for int_idx in int_idxs):
+        # We don't have a view-only broadcasting operation
+        # Explicitly broadcasting the indices can incur a memory / copy overhead
         return None
 
-    raveled_int_idx = int_idx.ravel()
-    new_idxs = list(idxs)
-    new_idxs[int_idx_pos] = raveled_int_idx
-    raveled_subtensor = x[tuple(new_idxs)]
-
-    # Reshape into correct shape
-    # Because we only allow one advanced indexing, the output dimension corresponding to the raveled integer indexing
-    # must match the input position. If there were multiple advanced indexes, this could have been forcefully moved to the front
-    raveled_shape = raveled_subtensor.shape
-    unraveled_shape = (
-        *raveled_shape[:int_idx_pos],
-        *int_idx.shape,
-        *raveled_shape[int_idx_pos + 1 :],
-    )
-    new_out = raveled_subtensor.reshape(unraveled_shape)
+    int_idxs_ndim = len(first_int_idx_bcast)
+    if (
+        int_idxs_ndim == 0
+    ):  # This should be a basic indexing operation, rewrite elsewhere
+        return None
+
+    int_idxs_need_raveling = int_idxs_ndim > 1
+    if not (int_idxs_need_raveling or non_consecutive_adv_indexing):
+        # Numba or our dispatch natively supports consecutive vector indices, nothing needs to be done
+        return None
+
+    # Reorder non-consecutive indices
+    if non_consecutive_adv_indexing:
+        assert not is_inc_subtensor  # Sanity check that we got out if this was the case
+        # This case works as if all the advanced indices were on the front
+        transposition = list(int_idxs_pos) + [
+            i for i in range(len(idxs)) if i not in int_idxs_pos
+        ]
+        idxs = tuple(idxs[a] for a in transposition)
+        x = x.transpose(transposition)
+        first_int_idx_pos = 0
+        del int_idxs_pos  # Make sure they are not wrongly used
+
+    # Ravel multidimensional indices
+    if int_idxs_need_raveling:
+        idxs = list(idxs)
+        for idx_pos, int_idx in enumerate(int_idxs, start=first_int_idx_pos):
+            idxs[idx_pos] = int_idx.ravel()
+
+    # Index with reordered and/or raveled indices
+    new_subtensor = x[tuple(idxs)]
+
+    if is_inc_subtensor:
+        y_shape = tuple(y.shape)
+        y_raveled_shape = (
+            *y_shape[:first_int_idx_pos],
+            -1,
+            *y_shape[first_int_idx_pos + int_idxs_ndim :],
+        )
+        y_raveled = y.reshape(y_raveled_shape)
+
+        new_out = inc_subtensor(
+            new_subtensor,
+            y_raveled,
+            set_instead_of_inc=op.set_instead_of_inc,
+            ignore_duplicates=op.ignore_duplicates,
+            inplace=op.inplace,
+        )
+
+    else:
+        # Unravel advanced indexing dimensions
+        raveled_shape = tuple(new_subtensor.shape)
+        unraveled_shape = (
+            *raveled_shape[:first_int_idx_pos],
+            *first_int_idx.shape,
+            *raveled_shape[first_int_idx_pos + 1 :],
+        )
+        new_out = new_subtensor.reshape(unraveled_shape)
+
     return [copy_stack_trace(node.outputs[0], new_out)]
 
 
@@ -2089,10 +2163,12 @@ optdb["specialize"].register(
     ravel_multidimensional_bool_idx.__name__,
     ravel_multidimensional_bool_idx,
     "numba",
+    use_db_name_as_tag=False,  # Not included if only "specialize" is requested
 )
 
 optdb["specialize"].register(
     ravel_multidimensional_int_idx.__name__,
     ravel_multidimensional_int_idx,
     "numba",
+    use_db_name_as_tag=False,  # Not included if only "specialize" is requested
 )
diff --git a/pytensor/tensor/slinalg.py b/pytensor/tensor/slinalg.py
index 94973810f..25ee69a07 100644
--- a/pytensor/tensor/slinalg.py
+++ b/pytensor/tensor/slinalg.py
@@ -296,13 +296,12 @@ class SolveBase(Op):
         # We need to return (dC/d[inv(A)], dC/db)
         c_bar = output_gradients[0]
 
-        trans_solve_op = type(self)(
-            **{
-                k: (not getattr(self, k) if k == "lower" else getattr(self, k))
-                for k in self.__props__
-            }
-        )
-        b_bar = trans_solve_op(A.T, c_bar)
+        props_dict = self._props_dict()
+        props_dict["lower"] = not self.lower
+
+        solve_op = type(self)(**props_dict)
+
+        b_bar = solve_op(A.T, c_bar)
         # force outer product if vector second input
         A_bar = -ptm.outer(b_bar, c) if c.ndim == 1 else -b_bar.dot(c.T)
 
@@ -385,7 +384,6 @@ class SolveTriangular(SolveBase):
     """Solve a system of linear equations."""
 
     __props__ = (
-        "trans",
         "unit_diagonal",
         "lower",
         "check_finite",
@@ -393,11 +391,10 @@ class SolveTriangular(SolveBase):
         "overwrite_b",
     )
 
-    def __init__(self, *, trans=0, unit_diagonal=False, **kwargs):
+    def __init__(self, *, unit_diagonal=False, **kwargs):
         if kwargs.get("overwrite_a", False):
             raise ValueError("overwrite_a is not supported for SolverTriangulare")
         super().__init__(**kwargs)
-        self.trans = trans
         self.unit_diagonal = unit_diagonal
 
     def perform(self, node, inputs, outputs):
@@ -406,7 +403,7 @@ class SolveTriangular(SolveBase):
             A,
             b,
             lower=self.lower,
-            trans=self.trans,
+            trans=0,
             unit_diagonal=self.unit_diagonal,
             check_finite=self.check_finite,
             overwrite_b=self.overwrite_b,
@@ -445,9 +442,9 @@ def solve_triangular(
 
     Parameters
     ----------
-    a
+    a: TensorVariable
         Square input data
-    b
+    b: TensorVariable
         Input data for the right hand side.
     lower : bool, optional
         Use only data contained in the lower triangle of `a`. Default is to use upper triangle.
@@ -468,10 +465,17 @@ def solve_triangular(
         This will influence how batched dimensions are interpreted.
     """
     b_ndim = _default_b_ndim(b, b_ndim)
+
+    if trans in [1, "T", True]:
+        a = a.mT
+        lower = not lower
+    if trans in [2, "C"]:
+        a = a.conj().mT
+        lower = not lower
+
     ret = Blockwise(
         SolveTriangular(
             lower=lower,
-            trans=trans,
             unit_diagonal=unit_diagonal,
             check_finite=check_finite,
             b_ndim=b_ndim,
@@ -534,6 +538,7 @@ def solve(
     *,
     assume_a="gen",
     lower=False,
+    transposed=False,
     check_finite=True,
     b_ndim: int | None = None,
 ):
@@ -564,8 +569,10 @@ def solve(
     b : (..., N, NRHS) array_like
         Input data for the right hand side.
     lower : bool, optional
-        If True, only the data contained in the lower triangle of `a`. Default
+        If True, use only the data contained in the lower triangle of `a`. Default
         is to use upper triangle. (ignored for ``'gen'``)
+    transposed: bool, optional
+        If True, solves the system A^T x = b. Default is False.
     check_finite : bool, optional
         Whether to check that the input matrices contain only finite numbers.
         Disabling may give a performance gain, but may result in problems
@@ -577,6 +584,11 @@ def solve(
         This will influence how batched dimensions are interpreted.
     """
     b_ndim = _default_b_ndim(b, b_ndim)
+
+    if transposed:
+        a = a.mT
+        lower = not lower
+
     return Blockwise(
         Solve(
             lower=lower,
diff --git a/pytensor/tensor/subtensor.py b/pytensor/tensor/subtensor.py
index 3a2304eb7..3de4f4106 100644
--- a/pytensor/tensor/subtensor.py
+++ b/pytensor/tensor/subtensor.py
@@ -1,5 +1,6 @@
 import logging
 import sys
+import warnings
 from collections.abc import Callable, Iterable
 from itertools import chain, groupby
 from textwrap import dedent
@@ -59,6 +60,7 @@ from pytensor.tensor.type import (
     zscalar,
 )
 from pytensor.tensor.type_other import (
+    MakeSlice,
     NoneConst,
     NoneTypeT,
     SliceConstant,
@@ -527,11 +529,20 @@ def basic_shape(shape, indices):
         if isinstance(idx, slice):
             res_shape += (slice_len(idx, n),)
         elif isinstance(getattr(idx, "type", None), SliceType):
-            if idx.owner:
-                idx_inputs = idx.owner.inputs
+            if idx.owner is None:
+                if not isinstance(idx, Constant):
+                    # This is an input slice, we can't reason symbolically on it.
+                    # We don't even know if we will get None entries or integers
+                    res_shape += (None,)
+                    continue
+                else:
+                    sl: slice = idx.data
+                    slice_inputs = (sl.start, sl.stop, sl.step)
+            elif isinstance(idx.owner.op, MakeSlice):
+                slice_inputs = idx.owner.inputs
             else:
-                idx_inputs = (None,)
-            res_shape += (slice_len(slice(*idx_inputs), n),)
+                raise ValueError(f"Unexpected Slice producing Op {idx.owner.op}")
+            res_shape += (slice_len(slice(*slice_inputs), n),)
         elif idx is None:
             res_shape += (ps.ScalarConstant(ps.int64, 1),)
         elif isinstance(getattr(idx, "type", None), NoneTypeT):
@@ -570,8 +581,8 @@ def group_indices(indices):
     return idx_groups
 
 
-def _non_contiguous_adv_indexing(indices) -> bool:
-    """Check if the advanced indexing is non-contiguous (i.e., split by basic indexing)."""
+def _non_consecutive_adv_indexing(indices) -> bool:
+    """Check if the advanced indexing is non-consecutive (i.e., split by basic indexing)."""
     idx_groups = group_indices(indices)
     # This means that there are at least two groups of advanced indexing separated by basic indexing
     return len(idx_groups) > 3 or (len(idx_groups) == 3 and not idx_groups[0][0])
@@ -601,7 +612,7 @@ def indexed_result_shape(array_shape, indices, indices_are_shapes=False):
     remaining_dims = range(pytensor.tensor.basic.get_vector_length(array_shape))
     idx_groups = group_indices(indices)
 
-    if _non_contiguous_adv_indexing(indices):
+    if _non_consecutive_adv_indexing(indices):
         # In this case NumPy places the advanced index groups in the front of the array
         # https://numpy.org/devdocs/user/basics.indexing.html#combining-advanced-and-basic-indexing
         idx_groups = sorted(idx_groups, key=lambda x: x[0])
@@ -2728,6 +2739,11 @@ class AdvancedSubtensor(Op):
         res_shape = list(
             indexed_result_shape(ishapes[0], index_shapes, indices_are_shapes=True)
         )
+        for i, res_dim_length in enumerate(res_shape):
+            if res_dim_length is None:
+                # This can happen when we have a Slice provided by the user (not a constant nor the result of MakeSlice)
+                # We must compute the Op to find its shape
+                res_shape[i] = Shape_i(i)(node.out)
 
         adv_indices = [idx for idx in indices if not is_basic_idx(idx)]
         bool_indices = [idx for idx in adv_indices if is_bool_index(idx)]
@@ -2781,10 +2797,17 @@ class AdvancedSubtensor(Op):
 
     @staticmethod
     def non_contiguous_adv_indexing(node: Apply) -> bool:
+        warnings.warn(
+            "Method was renamed to `non_consecutive_adv_indexing`", FutureWarning
+        )
+        return AdvancedSubtensor.non_consecutive_adv_indexing(node)
+
+    @staticmethod
+    def non_consecutive_adv_indexing(node: Apply) -> bool:
         """
-        Check if the advanced indexing is non-contiguous (i.e. interrupted by basic indexing).
+        Check if the advanced indexing is non-consecutive (i.e. interrupted by basic indexing).
 
-        This function checks if the advanced indexing is non-contiguous,
+        This function checks if the advanced indexing is non-consecutive,
         in which case the advanced index dimensions are placed on the left of the
         output array, regardless of their opriginal position.
 
@@ -2799,10 +2822,10 @@ class AdvancedSubtensor(Op):
         Returns
         -------
         bool
-            True if the advanced indexing is non-contiguous, False otherwise.
+            True if the advanced indexing is non-consecutive, False otherwise.
         """
         _, *idxs = node.inputs
-        return _non_contiguous_adv_indexing(idxs)
+        return _non_consecutive_adv_indexing(idxs)
 
 
 advanced_subtensor = AdvancedSubtensor()
@@ -2820,7 +2843,7 @@ def vectorize_advanced_subtensor(op: AdvancedSubtensor, node, *batch_inputs):
         if isinstance(batch_idx, TensorVariable)
     )
 
-    if idxs_are_batched or (x_is_batched and op.non_contiguous_adv_indexing(node)):
+    if idxs_are_batched or (x_is_batched and op.non_consecutive_adv_indexing(node)):
         # Fallback to Blockwise if idxs are batched or if we have non contiguous advanced indexing
         # which would put the indexed results to the left of the batch dimensions!
         # TODO: Not all cases must be handled by Blockwise, but the logic is complex
@@ -2829,7 +2852,7 @@ def vectorize_advanced_subtensor(op: AdvancedSubtensor, node, *batch_inputs):
         # TODO: Implement these internally, so Blockwise is always a safe fallback
         if any(not isinstance(idx, TensorVariable) for idx in idxs):
             raise NotImplementedError(
-                "Vectorized AdvancedSubtensor with batched indexes or non-contiguous advanced indexing "
+                "Vectorized AdvancedSubtensor with batched indexes or non-consecutive advanced indexing "
                 "and slices or newaxis is currently not supported."
             )
         else:
@@ -2939,10 +2962,17 @@ class AdvancedIncSubtensor(Op):
 
     @staticmethod
     def non_contiguous_adv_indexing(node: Apply) -> bool:
+        warnings.warn(
+            "Method was renamed to `non_consecutive_adv_indexing`", FutureWarning
+        )
+        return AdvancedIncSubtensor.non_consecutive_adv_indexing(node)
+
+    @staticmethod
+    def non_consecutive_adv_indexing(node: Apply) -> bool:
         """
-        Check if the advanced indexing is non-contiguous (i.e. interrupted by basic indexing).
+        Check if the advanced indexing is non-consecutive (i.e. interrupted by basic indexing).
 
-        This function checks if the advanced indexing is non-contiguous,
+        This function checks if the advanced indexing is non-consecutive,
         in which case the advanced index dimensions are placed on the left of the
         output array, regardless of their opriginal position.
 
@@ -2957,10 +2987,10 @@ class AdvancedIncSubtensor(Op):
         Returns
         -------
         bool
-            True if the advanced indexing is non-contiguous, False otherwise.
+            True if the advanced indexing is non-consecutive, False otherwise.
         """
         _, _, *idxs = node.inputs
-        return _non_contiguous_adv_indexing(idxs)
+        return _non_consecutive_adv_indexing(idxs)
 
 
 advanced_inc_subtensor = AdvancedIncSubtensor()
diff --git a/pytensor/tensor/variable.py b/pytensor/tensor/variable.py
index bad700de8..7e5d22528 100644
--- a/pytensor/tensor/variable.py
+++ b/pytensor/tensor/variable.py
@@ -346,7 +346,7 @@ class _tensor_py_operators:
         DimShuffle
 
         """
-        if (len(pattern) == 1) and (isinstance(pattern[0], list | tuple)):
+        if (len(pattern) == 1) and (isinstance(pattern[0], list | tuple | np.ndarray)):
             pattern = pattern[0]
         ds_op = pt.elemwise.DimShuffle(input_ndim=self.type.ndim, new_order=pattern)
         return ds_op(self)
diff --git a/scripts/mypy-failing.txt b/scripts/mypy-failing.txt
index a7cb4a182..99dd26a26 100644
--- a/scripts/mypy-failing.txt
+++ b/scripts/mypy-failing.txt
@@ -11,7 +11,6 @@ pytensor/link/numba/dispatch/scan.py
 pytensor/printing.py
 pytensor/raise_op.py
 pytensor/sparse/basic.py
-pytensor/sparse/type.py
 pytensor/tensor/basic.py
 pytensor/tensor/blas_c.py
 pytensor/tensor/blas_headers.py
diff --git a/scripts/run_mypy.py b/scripts/run_mypy.py
index c2e87560c..34cc81064 100644
--- a/scripts/run_mypy.py
+++ b/scripts/run_mypy.py
@@ -142,7 +142,13 @@ if __name__ == "__main__":
         print(*missing, sep="\n")
         sys.exit(1)
     cp = subprocess.run(
-        ["mypy", "--show-error-codes", "pytensor"],
+        [
+            "mypy",
+            "--show-error-codes",
+            "--disable-error-code",
+            "annotation-unchecked",
+            "pytensor",
+        ],
         capture_output=True,
     )
     output = cp.stdout.decode()


Test patch:
diff --git a/tests/link/jax/test_pad.py b/tests/link/jax/test_pad.py
index 8ecb460ac..13d71be9a 100644
--- a/tests/link/jax/test_pad.py
+++ b/tests/link/jax/test_pad.py
@@ -1,5 +1,6 @@
 import numpy as np
 import pytest
+from packaging import version
 
 import pytensor.tensor as pt
 from pytensor import config
@@ -16,7 +17,14 @@ RTOL = ATOL = 1e-6 if floatX.endswith("64") else 1e-3
     "mode, kwargs",
     [
         ("constant", {"constant_values": 0}),
-        ("constant", {"constant_values": (1, 2)}),
+        pytest.param(
+            "constant",
+            {"constant_values": (1, 2)},
+            marks=pytest.mark.skipif(
+                version.parse(jax.__version__) > version.parse("0.4.35"),
+                reason="Bug in JAX: https://github.com/jax-ml/jax/issues/26888",
+            ),
+        ),
         ("edge", {}),
         ("linear_ramp", {"end_values": 0}),
         ("linear_ramp", {"end_values": (1, 2)}),
diff --git a/tests/link/jax/test_slinalg.py b/tests/link/jax/test_slinalg.py
index 2656b0fd0..c446437dd 100644
--- a/tests/link/jax/test_slinalg.py
+++ b/tests/link/jax/test_slinalg.py
@@ -5,6 +5,7 @@ import numpy as np
 import pytest
 
 import pytensor.tensor as pt
+import tests.unittest_tools as utt
 from pytensor.configdefaults import config
 from pytensor.tensor import nlinalg as pt_nlinalg
 from pytensor.tensor import slinalg as pt_slinalg
@@ -103,28 +104,41 @@ def test_jax_basic():
     )
 
 
-@pytest.mark.parametrize("check_finite", [False, True])
-@pytest.mark.parametrize("lower", [False, True])
-@pytest.mark.parametrize("trans", [0, 1, 2])
-def test_jax_SolveTriangular(trans, lower, check_finite):
-    x = matrix("x")
-    b = vector("b")
+def test_jax_solve():
+    rng = np.random.default_rng(utt.fetch_seed())
+
+    A = pt.tensor("A", shape=(5, 5))
+    b = pt.tensor("B", shape=(5, 5))
+
+    out = pt_slinalg.solve(A, b, lower=False, transposed=False)
+
+    A_val = rng.normal(size=(5, 5)).astype(config.floatX)
+    b_val = rng.normal(size=(5, 5)).astype(config.floatX)
 
-    out = pt_slinalg.solve_triangular(
-        x,
-        b,
-        trans=trans,
-        lower=lower,
-        check_finite=check_finite,
-    )
     compare_jax_and_py(
-        [x, b],
+        [A, b],
         [out],
-        [
-            np.eye(10).astype(config.floatX),
-            np.arange(10).astype(config.floatX),
-        ],
+        [A_val, b_val],
+    )
+
+
+def test_jax_SolveTriangular():
+    rng = np.random.default_rng(utt.fetch_seed())
+
+    A = pt.tensor("A", shape=(5, 5))
+    b = pt.tensor("B", shape=(5, 5))
+
+    A_val = rng.normal(size=(5, 5)).astype(config.floatX)
+    b_val = rng.normal(size=(5, 5)).astype(config.floatX)
+
+    out = pt_slinalg.solve_triangular(
+        A,
+        b,
+        trans=0,
+        lower=True,
+        unit_diagonal=False,
     )
+    compare_jax_and_py([A, b], [out], [A_val, b_val])
 
 
 def test_jax_block_diag():
diff --git a/tests/link/numba/test_slinalg.py b/tests/link/numba/test_slinalg.py
index 5caeb8bef..defbcf6c8 100644
--- a/tests/link/numba/test_slinalg.py
+++ b/tests/link/numba/test_slinalg.py
@@ -5,7 +5,6 @@ from typing import Literal
 import numpy as np
 import pytest
 from numpy.testing import assert_allclose
-from scipy import linalg as scipy_linalg
 
 import pytensor
 import pytensor.tensor as pt
@@ -26,9 +25,9 @@ def transpose_func(x, trans):
     if trans == 0:
         return x
     if trans == 1:
-        return x.conj().T
-    if trans == 2:
         return x.T
+    if trans == 2:
+        return x.conj().T
 
 
 @pytest.mark.parametrize(
@@ -59,18 +58,18 @@ def test_solve_triangular(b_shape: tuple[int], lower, trans, unit_diag, is_compl
 
     def A_func(x):
         x = x @ x.conj().T
-        x_tri = scipy_linalg.cholesky(x, lower=lower).astype(dtype)
+        x_tri = pt.linalg.cholesky(x, lower=lower).astype(dtype)
 
         if unit_diag:
-            x_tri[np.diag_indices_from(x_tri)] = 1.0
+            x_tri = pt.fill_diagonal(x_tri, 1.0)
 
-        return x_tri.astype(dtype)
+        return x_tri
 
     solve_op = partial(
         pt.linalg.solve_triangular, lower=lower, trans=trans, unit_diagonal=unit_diag
     )
 
-    X = solve_op(A, b)
+    X = solve_op(A_func(A), b)
     f = pytensor.function([A, b], X, mode="NUMBA")
 
     A_val = np.random.normal(size=(5, 5))
@@ -80,20 +79,20 @@ def test_solve_triangular(b_shape: tuple[int], lower, trans, unit_diag, is_compl
         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_func(A_val), b_val)
-
-    test_input = transpose_func(A_func(A_val), trans)
-
-    ATOL = 1e-8 if floatX.endswith("64") else 1e-4
-    RTOL = 1e-8 if floatX.endswith("64") else 1e-4
-
-    np.testing.assert_allclose(test_input @ X_np, b_val, atol=ATOL, rtol=RTOL)
+    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_func(A_val), b_val],
+        [A_val, b_val],
     )
 
 
@@ -145,7 +144,6 @@ def test_solve_triangular_overwrite_b_correct(overwrite_b):
     b_test_nb = b_test_py.copy(order="F")
 
     op = SolveTriangular(
-        trans=0,
         unit_diagonal=False,
         lower=False,
         check_finite=True,
diff --git a/tests/link/numba/test_subtensor.py b/tests/link/numba/test_subtensor.py
index 8b95de34b..675afdc99 100644
--- a/tests/link/numba/test_subtensor.py
+++ b/tests/link/numba/test_subtensor.py
@@ -81,11 +81,6 @@ def test_AdvancedSubtensor1_out_of_bounds():
             (np.array([True, False, False])),
             False,
         ),
-        (
-            pt.as_tensor(np.arange(3 * 4 * 5).reshape((3, 4, 5))),
-            ([1, 2], [2, 3]),
-            False,
-        ),
         # Single multidimensional indexing (supported after specialization rewrites)
         (
             as_tensor(np.arange(3 * 3).reshape((3, 3))),
@@ -117,6 +112,12 @@ def test_AdvancedSubtensor1_out_of_bounds():
             (slice(2, None), np.eye(3).astype(bool)),
             False,
         ),
+        # Multiple vector indexing (supported by our dispatcher)
+        (
+            pt.as_tensor(np.arange(3 * 4 * 5).reshape((3, 4, 5))),
+            ([1, 2], [2, 3]),
+            False,
+        ),
         (
             as_tensor(np.arange(3 * 4 * 5).reshape((3, 4, 5))),
             (slice(None), [1, 2], [3, 4]),
@@ -127,18 +128,35 @@ def test_AdvancedSubtensor1_out_of_bounds():
             ([1, 2], [3, 4], [5, 6]),
             False,
         ),
-        # Non-contiguous vector indexing, only supported in obj mode
+        # Non-consecutive vector indexing, supported by our dispatcher after rewriting
         (
             as_tensor(np.arange(3 * 4 * 5).reshape((3, 4, 5))),
             ([1, 2], slice(None), [3, 4]),
-            True,
+            False,
+        ),
+        # Multiple multidimensional integer indexing (supported by our dispatcher)
+        (
+            as_tensor(np.arange(3 * 4 * 5).reshape((3, 4, 5))),
+            ([[1, 2], [2, 1]], [[0, 0], [0, 0]]),
+            False,
+        ),
+        (
+            as_tensor(np.arange(2 * 3 * 4 * 5).reshape((2, 3, 4, 5))),
+            (slice(None), [[1, 2], [2, 1]], slice(None), [[0, 0], [0, 0]]),
+            False,
         ),
-        # >1d vector indexing, only supported in obj mode
+        # Multiple multidimensional indexing with broadcasting, only supported in obj mode
         (
             as_tensor(np.arange(3 * 4 * 5).reshape((3, 4, 5))),
             ([[1, 2], [2, 1]], [0, 0]),
             True,
         ),
+        # multiple multidimensional integer indexing mixed with basic indexing, only supported in obj mode
+        (
+            as_tensor(np.arange(3 * 4 * 5).reshape((3, 4, 5))),
+            ([[1, 2], [2, 1]], slice(1, None), [[0, 0], [0, 0]]),
+            True,
+        ),
     ],
 )
 @pytest.mark.filterwarnings("error")  # Raise if we did not expect objmode to be needed
@@ -297,7 +315,7 @@ def test_AdvancedIncSubtensor1(x, y, indices):
         (
             np.arange(3 * 4 * 5).reshape((3, 4, 5)),
             -np.arange(4 * 5).reshape(4, 5),
-            (0, [1, 2, 2, 3]),  # Broadcasted vector index
+            (0, [1, 2, 2, 3]),  # Broadcasted vector index with repeated values
             True,
             False,
             True,
@@ -305,7 +323,7 @@ def test_AdvancedIncSubtensor1(x, y, indices):
         (
             np.arange(3 * 4 * 5).reshape((3, 4, 5)),
             np.array([-99]),  # Broadcasted value
-            (0, [1, 2, 2, 3]),  # Broadcasted vector index
+            (0, [1, 2, 2, 3]),  # Broadcasted vector index with repeated values
             True,
             False,
             True,
@@ -380,7 +398,7 @@ def test_AdvancedIncSubtensor1(x, y, indices):
         (
             np.arange(3 * 4 * 5).reshape((3, 4, 5)),
             rng.poisson(size=(2, 4)),
-            ([1, 2], slice(None), [3, 4]),  # Non-contiguous vector indices
+            ([1, 2], slice(None), [3, 4]),  # Non-consecutive vector indices
             False,
             True,
             True,
@@ -400,15 +418,23 @@ def test_AdvancedIncSubtensor1(x, y, indices):
         (
             np.arange(5),
             rng.poisson(size=(2, 2)),
-            ([[1, 2], [2, 3]]),  # matrix indices
+            ([[1, 2], [2, 3]]),  # matrix index
+            False,
+            False,
+            False,
+        ),
+        (
+            np.arange(3 * 5).reshape((3, 5)),
+            rng.poisson(size=(2, 2, 2)),
+            (slice(1, 3), [[1, 2], [2, 3]]),  # matrix index, mixed with basic index
+            False,
+            False,
             False,
-            False,  # Gets converted to AdvancedIncSubtensor1
-            True,  # This is actually supported with the default `ignore_duplicates=False`
         ),
         (
             np.arange(3 * 5).reshape((3, 5)),
-            rng.poisson(size=(1, 2, 2)),
-            (slice(1, 3), [[1, 2], [2, 3]]),  # matrix indices, mixed with basic index
+            rng.poisson(size=(1, 2, 2)),  # Same as before, but Y broadcasts
+            (slice(1, 3), [[1, 2], [2, 3]]),
             False,
             True,
             True,
@@ -421,6 +447,14 @@ def test_AdvancedIncSubtensor1(x, y, indices):
             False,
             False,
         ),
+        (
+            np.arange(3 * 4 * 5).reshape((3, 4, 5)),
+            rng.poisson(size=(3, 2, 2)),
+            (slice(None), [[1, 2], [2, 1]], [[2, 3], [0, 0]]),  # 2 matrix indices
+            False,
+            False,
+            False,
+        ),
     ],
 )
 @pytest.mark.parametrize("inplace", (False, True))
diff --git a/tests/tensor/test_einsum.py b/tests/tensor/test_einsum.py
index ba8e35451..951e9a0c5 100644
--- a/tests/tensor/test_einsum.py
+++ b/tests/tensor/test_einsum.py
@@ -8,6 +8,7 @@ import pytensor
 from pytensor import Mode, config, function
 from pytensor.graph import FunctionGraph
 from pytensor.graph.op import HasInnerGraph
+from pytensor.tensor import matrix
 from pytensor.tensor.basic import moveaxis
 from pytensor.tensor.blockwise import Blockwise
 from pytensor.tensor.einsum import _delta, _general_dot, _iota, einsum
@@ -281,3 +282,15 @@ def test_threeway_mul(static_length):
         out.eval({x: x_test, y: y_test, z: z_test}),
         np.full((3,), fill_value=6),
     )
+
+
+def test_repeated_inputs():
+    x = matrix("x")
+    out_repeated = einsum("ij,ij->i", x, x)
+    out_copy = einsum("ij,ij->i", x, x.copy())
+
+    x_test = np.array([[1, 2], [3, 4]]).astype(x.dtype)
+
+    np.testing.assert_allclose(
+        out_repeated.eval({x: x_test}), out_copy.eval({x: x_test})
+    )
diff --git a/tests/tensor/test_elemwise.py b/tests/tensor/test_elemwise.py
index d5aac0113..77d41a03c 100644
--- a/tests/tensor/test_elemwise.py
+++ b/tests/tensor/test_elemwise.py
@@ -185,14 +185,13 @@ class TestDimShuffle(unittest_tools.InferShapeTester):
         # as the broadcasted value; that way, we'll be able to tell that we're getting
         # junk data from a poorly constructed array view.
         x_val = np.broadcast_to(2039, (5000,))
-        expected_x_val = x_val[None]
         for i in range(1):
             inputs[0].storage[0] = x_val
             thunk()
             # Make sure it's a view of the original data
             assert np.shares_memory(x_val, outputs[0].storage[0])
             # Confirm the right strides
-            assert outputs[0].storage[0].strides == expected_x_val.strides
+            assert outputs[0].storage[0].strides[-1] == 0
             # Confirm the broadcasted value in the output
             assert np.array_equiv(outputs[0].storage[0], 2039)
 
diff --git a/tests/tensor/test_math.py b/tests/tensor/test_math.py
index 9ab4fd104..38207d0f5 100644
--- a/tests/tensor/test_math.py
+++ b/tests/tensor/test_math.py
@@ -89,6 +89,7 @@ from pytensor.tensor.math import (
     logaddexp,
     logsumexp,
     matmul,
+    matvec,
     max,
     max_and_argmax,
     maximum,
@@ -123,6 +124,8 @@ from pytensor.tensor.math import (
     true_div,
     trunc,
     var,
+    vecdot,
+    vecmat,
 )
 from pytensor.tensor.math import sum as pt_sum
 from pytensor.tensor.type import (
@@ -2076,6 +2079,71 @@ class TestDot:
                             assert is_super_shape(y, g)
 
 
+def test_matrix_vector_ops():
+    """Test vecdot, matvec, and vecmat helper functions."""
+    rng = np.random.default_rng(seed=utt.fetch_seed())
+
+    # Create test data with batch dimension (2)
+    batch_size = 2
+    dim_k = 4  # Common dimension
+    dim_m = 3  # Matrix rows
+    dim_n = 5  # Matrix columns
+
+    # Create input tensors with appropriate shapes
+    # For matvec: x1(b,m,k) @ x2(b,k) -> out(b,m)
+    # For vecmat: x1(b,k) @ x2(b,k,n) -> out(b,n)
+
+    # Create test values using config.floatX to match PyTensor's default dtype
+    mat_mk_val = random(batch_size, dim_m, dim_k, rng=rng).astype(config.floatX)
+    mat_kn_val = random(batch_size, dim_k, dim_n, rng=rng).astype(config.floatX)
+    vec_k_val = random(batch_size, dim_k, rng=rng).astype(config.floatX)
+
+    # Create tensor variables with matching dtype
+    mat_mk = tensor(
+        name="mat_mk", shape=(batch_size, dim_m, dim_k), dtype=config.floatX
+    )
+    mat_kn = tensor(
+        name="mat_kn", shape=(batch_size, dim_k, dim_n), dtype=config.floatX
+    )
+    vec_k = tensor(name="vec_k", shape=(batch_size, dim_k), dtype=config.floatX)
+
+    # Test 1: vecdot with matching dimensions
+    vecdot_out = vecdot(vec_k, vec_k, dtype="int32")
+    vecdot_fn = function([vec_k], vecdot_out)
+    result = vecdot_fn(vec_k_val)
+
+    # Check dtype
+    assert result.dtype == np.int32
+
+    # Calculate expected manually
+    expected_vecdot = np.zeros((batch_size,), dtype=np.int32)
+    for i in range(batch_size):
+        expected_vecdot[i] = np.sum(vec_k_val[i] * vec_k_val[i])
+    np.testing.assert_allclose(result, expected_vecdot)
+
+    # Test 2: matvec - matrix-vector product
+    matvec_out = matvec(mat_mk, vec_k)
+    matvec_fn = function([mat_mk, vec_k], matvec_out)
+    result_matvec = matvec_fn(mat_mk_val, vec_k_val)
+
+    # Calculate expected manually
+    expected_matvec = np.zeros((batch_size, dim_m), dtype=config.floatX)
+    for i in range(batch_size):
+        expected_matvec[i] = np.dot(mat_mk_val[i], vec_k_val[i])
+    np.testing.assert_allclose(result_matvec, expected_matvec)
+
+    # Test 3: vecmat - vector-matrix product
+    vecmat_out = vecmat(vec_k, mat_kn)
+    vecmat_fn = function([vec_k, mat_kn], vecmat_out)
+    result_vecmat = vecmat_fn(vec_k_val, mat_kn_val)
+
+    # Calculate expected manually
+    expected_vecmat = np.zeros((batch_size, dim_n), dtype=config.floatX)
+    for i in range(batch_size):
+        expected_vecmat[i] = np.dot(vec_k_val[i], mat_kn_val[i])
+    np.testing.assert_allclose(result_vecmat, expected_vecmat)
+
+
 class TestTensordot:
     def TensorDot(self, axes):
         # Since tensordot is no longer an op, mimic the old op signature
diff --git a/tests/tensor/test_slinalg.py b/tests/tensor/test_slinalg.py
index 34f1396f4..f1a6b0fe5 100644
--- a/tests/tensor/test_slinalg.py
+++ b/tests/tensor/test_slinalg.py
@@ -214,7 +214,38 @@ def test_solve_raises_on_invalid_A():
         Solve(assume_a="test", b_ndim=2)
 
 
+solve_test_cases = [
+    ("gen", False, False),
+    ("gen", False, True),
+    ("sym", False, False),
+    ("sym", True, False),
+    ("sym", True, True),
+    ("pos", False, False),
+    ("pos", True, False),
+    ("pos", True, True),
+]
+solve_test_ids = [
+    f'{assume_a}_{"lower" if lower else "upper"}_{"A^T" if transposed else "A"}'
+    for assume_a, lower, transposed in solve_test_cases
+]
+
+
 class TestSolve(utt.InferShapeTester):
+    @staticmethod
+    def A_func(x, assume_a):
+        if assume_a == "pos":
+            return x @ x.T
+        elif assume_a == "sym":
+            return (x + x.T) / 2
+        else:
+            return x
+
+    @staticmethod
+    def T(x, transposed):
+        if transposed:
+            return x.T
+        return x
+
     @pytest.mark.parametrize("b_shape", [(5, 1), (5,)])
     def test_infer_shape(self, b_shape):
         rng = np.random.default_rng(utt.fetch_seed())
@@ -235,8 +266,12 @@ class TestSolve(utt.InferShapeTester):
     @pytest.mark.parametrize(
         "b_size", [(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_correctness(self, b_size: tuple[int], assume_a: str):
+    @pytest.mark.parametrize(
+        "assume_a, lower, transposed", solve_test_cases, ids=solve_test_ids
+    )
+    def test_solve_correctness(
+        self, b_size: tuple[int], assume_a: str, lower: bool, transposed: bool
+    ):
         rng = np.random.default_rng(utt.fetch_seed())
         A = pt.tensor("A", shape=(5, 5))
         b = pt.tensor("b", shape=b_size)
@@ -244,19 +279,18 @@ class TestSolve(utt.InferShapeTester):
         A_val = rng.normal(size=(5, 5)).astype(config.floatX)
         b_val = rng.normal(size=b_size).astype(config.floatX)
 
-        solve_op = functools.partial(solve, assume_a=assume_a, b_ndim=len(b_size))
+        A_func = functools.partial(self.A_func, assume_a=assume_a)
+        T = functools.partial(self.T, transposed=transposed)
 
-        def A_func(x):
-            if assume_a == "pos":
-                return x @ x.T
-            elif assume_a == "sym":
-                return (x + x.T) / 2
-            else:
-                return x
-
-        solve_input_val = A_func(A_val)
+        y = solve(
+            A_func(A),
+            b,
+            assume_a=assume_a,
+            lower=lower,
+            transposed=transposed,
+            b_ndim=len(b_size),
+        )
 
-        y = solve_op(A_func(A), b)
         solve_func = pytensor.function([A, b], y)
         X_np = solve_func(A_val.copy(), b_val.copy())
 
@@ -264,22 +298,34 @@ class TestSolve(utt.InferShapeTester):
         RTOL = 1e-8 if config.floatX.endswith("64") else 1e-4
 
         np.testing.assert_allclose(
-            scipy.linalg.solve(solve_input_val, b_val, assume_a=assume_a),
+            scipy.linalg.solve(
+                A_func(A_val),
+                b_val,
+                assume_a=assume_a,
+                transposed=transposed,
+                lower=lower,
+            ),
             X_np,
             atol=ATOL,
             rtol=RTOL,
         )
 
-        np.testing.assert_allclose(A_func(A_val) @ X_np, b_val, atol=ATOL, rtol=RTOL)
+        np.testing.assert_allclose(T(A_func(A_val)) @ X_np, b_val, atol=ATOL, rtol=RTOL)
 
     @pytest.mark.parametrize(
         "b_size", [(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.parametrize(
+        "assume_a, lower, transposed",
+        solve_test_cases,
+        ids=solve_test_ids,
+    )
     @pytest.mark.skipif(
         config.floatX == "float32", reason="Gradients not numerically stable in float32"
     )
-    def test_solve_gradient(self, b_size: tuple[int], assume_a: str):
+    def test_solve_gradient(
+        self, b_size: tuple[int], assume_a: str, lower: bool, transposed: bool
+    ):
         rng = np.random.default_rng(utt.fetch_seed())
 
         eps = 2e-8 if config.floatX == "float64" else None
@@ -287,15 +333,8 @@ class TestSolve(utt.InferShapeTester):
         A_val = rng.normal(size=(5, 5)).astype(config.floatX)
         b_val = rng.normal(size=b_size).astype(config.floatX)
 
-        def A_func(x):
-            if assume_a == "pos":
-                return x @ x.T
-            elif assume_a == "sym":
-                return (x + x.T) / 2
-            else:
-                return x
-
         solve_op = functools.partial(solve, assume_a=assume_a, b_ndim=len(b_size))
+        A_func = functools.partial(self.A_func, assume_a=assume_a)
 
         # To correctly check the gradients, we need to include a transformation from the space of unconstrained matrices
         # (A) to a valid input matrix for the given solver. This is done by the A_func function. If this isn't included,
@@ -307,11 +346,27 @@ class TestSolve(utt.InferShapeTester):
 
 
 class TestSolveTriangular(utt.InferShapeTester):
+    @staticmethod
+    def A_func(x, lower, unit_diagonal):
+        x = x @ x.T
+        x = pt.linalg.cholesky(x, lower=lower)
+        if unit_diagonal:
+            x = pt.fill_diagonal(x, 1)
+        return x
+
+    @staticmethod
+    def T(x, trans):
+        if trans == 1:
+            return x.T
+        elif trans == 2:
+            return x.conj().T
+        return x
+
     @pytest.mark.parametrize("b_shape", [(5, 1), (5,)])
     def test_infer_shape(self, b_shape):
         rng = np.random.default_rng(utt.fetch_seed())
         A = matrix()
-        b_val = np.asarray(rng.random(b_shape), dtype=config.floatX)
+        b_val = rng.random(b_shape).astype(config.floatX)
         b = pt.as_tensor_variable(b_val).type()
         self._compile_and_check(
             [A, b],
@@ -324,56 +379,78 @@ class TestSolveTriangular(utt.InferShapeTester):
             warn=False,
         )
 
+    @pytest.mark.parametrize(
+        "b_shape", [(5, 1), (5,), (5, 5)], ids=["b_col_vec", "b_vec", "b_matrix"]
+    )
     @pytest.mark.parametrize("lower", [True, False])
-    def test_correctness(self, lower):
+    @pytest.mark.parametrize("trans", [0, 1, 2])
+    @pytest.mark.parametrize("unit_diagonal", [True, False])
+    def test_correctness(self, b_shape: tuple[int], lower, trans, unit_diagonal):
         rng = np.random.default_rng(utt.fetch_seed())
+        A = pt.tensor("A", shape=(5, 5))
+        b = pt.tensor("b", shape=b_shape)
 
-        b_val = np.asarray(rng.random((5, 1)), dtype=config.floatX)
+        A_val = rng.random((5, 5)).astype(config.floatX)
+        b_val = rng.random(b_shape).astype(config.floatX)
 
-        A_val = np.asarray(rng.random((5, 5)), dtype=config.floatX)
-        A_val = np.dot(A_val.transpose(), A_val)
+        A_func = functools.partial(
+            self.A_func, lower=lower, unit_diagonal=unit_diagonal
+        )
 
-        C_val = scipy.linalg.cholesky(A_val, lower=lower)
+        x = solve_triangular(
+            A_func(A),
+            b,
+            lower=lower,
+            trans=trans,
+            unit_diagonal=unit_diagonal,
+            b_ndim=len(b_shape),
+        )
 
-        A = matrix()
-        b = matrix()
+        f = pytensor.function([A, b], x)
 
-        cholesky = Cholesky(lower=lower)
-        C = cholesky(A)
-        y_lower = solve_triangular(C, b, lower=lower)
-        lower_solve_func = pytensor.function([C, b], y_lower)
+        x_pt = f(A_val, b_val)
+        x_sp = scipy.linalg.solve_triangular(
+            A_func(A_val).eval(),
+            b_val,
+            lower=lower,
+            trans=trans,
+            unit_diagonal=unit_diagonal,
+        )
 
-        assert np.allclose(
-            scipy.linalg.solve_triangular(C_val, b_val, lower=lower),
-            lower_solve_func(C_val, b_val),
+        np.testing.assert_allclose(
+            x_pt,
+            x_sp,
+            atol=1e-8 if config.floatX == "float64" else 1e-4,
+            rtol=1e-8 if config.floatX == "float64" else 1e-4,
         )
 
     @pytest.mark.parametrize(
-        "m, n, lower",
-        [
-            (5, None, False),
-            (5, None, True),
-            (4, 2, False),
-            (4, 2, True),
-        ],
+        "b_shape", [(5, 1), (5,), (5, 5)], ids=["b_col_vec", "b_vec", "b_matrix"]
     )
-    def test_solve_grad(self, m, n, lower):
-        rng = np.random.default_rng(utt.fetch_seed())
+    @pytest.mark.parametrize("lower", [True, False])
+    @pytest.mark.parametrize("trans", [0, 1])
+    @pytest.mark.parametrize("unit_diagonal", [True, False])
+    def test_solve_triangular_grad(self, b_shape, lower, trans, unit_diagonal):
+        if config.floatX == "float32":
+            pytest.skip(reason="Not enough precision in float32 to get a good gradient")
 
-        # Ensure diagonal elements of `A` are relatively large to avoid
-        # numerical precision issues
-        A_val = (rng.normal(size=(m, m)) * 0.5 + np.eye(m)).astype(config.floatX)
+        rng = np.random.default_rng(utt.fetch_seed())
+        A_val = rng.normal(size=(5, 5)).astype(config.floatX)
+        b_val = rng.normal(size=b_shape).astype(config.floatX)
 
-        if n is None:
-            b_val = rng.normal(size=m).astype(config.floatX)
-        else:
-            b_val = rng.normal(size=(m, n)).astype(config.floatX)
+        A_func = functools.partial(
+            self.A_func, lower=lower, unit_diagonal=unit_diagonal
+        )
 
         eps = None
         if config.floatX == "float64":
             eps = 2e-8
 
-        solve_op = SolveTriangular(lower=lower, b_ndim=1 if n is None else 2)
+        def solve_op(A, b):
+            return solve_triangular(
+                A_func(A), b, lower=lower, trans=trans, unit_diagonal=unit_diagonal
+            )
+
         utt.verify_grad(solve_op, [A_val, b_val], 3, rng, eps=eps)
 
 
diff --git a/tests/tensor/test_subtensor.py b/tests/tensor/test_subtensor.py
index 3886a08f4..ebe07f494 100644
--- a/tests/tensor/test_subtensor.py
+++ b/tests/tensor/test_subtensor.py
@@ -15,6 +15,7 @@ from pytensor.compile.io import In
 from pytensor.compile.mode import Mode
 from pytensor.configdefaults import config
 from pytensor.gradient import grad
+from pytensor.graph import Constant
 from pytensor.graph.op import get_test_value
 from pytensor.graph.rewriting.utils import is_same_graph
 from pytensor.printing import pprint
@@ -37,6 +38,7 @@ from pytensor.tensor.subtensor import (
     advanced_inc_subtensor1,
     advanced_set_subtensor,
     advanced_set_subtensor1,
+    advanced_subtensor,
     advanced_subtensor1,
     as_index_literal,
     basic_shape,
@@ -2145,7 +2147,17 @@ class TestAdvancedSubtensor:
         slc = slicetype()
         f = pytensor.function([slc], var[slc], mode=self.mode)
         s = slice(1, 3)
-        f(s)
+        assert f(s).shape == (2, 3)
+
+        f_shape0 = pytensor.function([slc], var[slc].shape[0], mode=self.mode)
+        assert f_shape0(s) == 2
+
+        f_shape1 = pytensor.function([slc], var[slc].shape[1], mode=self.mode)
+        assert not any(
+            isinstance(node.op, AdvancedSubtensor)
+            for node in f_shape1.maker.fgraph.toposort()
+        )
+        assert f_shape1(s) == 3
 
     def test_adv_grouped(self):
         # Reported in https://github.com/Theano/Theano/issues/6152
@@ -2611,6 +2623,14 @@ class TestInferShape(utt.InferShapeTester):
             AdvancedSubtensor,
         )
 
+    def test_advanced_subtensor_constant_slice(self):
+        x = dmatrix("x")
+        constant_slice = pytensor.as_symbolic(slice(1, None, None))
+        assert isinstance(constant_slice, Constant)
+        adv_indices = ptb.constant(np.zeros((2, 3)), dtype="int")
+        y = advanced_subtensor(x, constant_slice, adv_indices)
+        assert tuple(y.shape.eval({x: np.zeros((10, 10))})) == (9, 2, 3)
+
 
 @config.change_flags(compute_test_value="raise")
 def test_basic_shape():
diff --git a/tests/tensor/test_variable.py b/tests/tensor/test_variable.py
index 57e47ce06..2c6f818c3 100644
--- a/tests/tensor/test_variable.py
+++ b/tests/tensor/test_variable.py
@@ -451,6 +451,21 @@ class TestTensorInstanceMethods:
         with pytest.raises(TypeError, match=msg):
             x[0] += 5
 
+    def test_transpose(self):
+        X, _ = self.vars
+        x, _ = self.vals
+
+        # Turn (2,2) -> (1,2)
+        X, x = X[1:, :], x[1:, :]
+
+        assert_array_equal(X.transpose(0, 1).eval({X: x}), x.transpose(0, 1))
+        assert_array_equal(X.transpose(1, 0).eval({X: x}), x.transpose(1, 0))
+
+        # Test handing in tuples, lists and np.arrays
+        equal_computations([X.transpose((1, 0))], [X.transpose(1, 0)])
+        equal_computations([X.transpose([1, 0])], [X.transpose(1, 0)])
+        equal_computations([X.transpose(np.array([1, 0]))], [X.transpose(1, 0)])
+
 
 def test_deprecated_import():
     with pytest.warns(