Fix GetItemList setting an incorrect static output shape

pytensor/sparse/basic.py

@@ -27,6 +27,7 @@
 from pytensor.tensor.basic import Split
 from pytensor.tensor.math import minimum
 from pytensor.tensor.shape import specify_broadcastable
+from pytensor.tensor.subtensor import slice_static_length
 from pytensor.tensor.type import TensorType, ivector, scalar, tensor, vector
 from pytensor.tensor.type import continuous_dtypes as tensor_continuous_dtypes
 from pytensor.tensor.type import discrete_dtypes as tensor_discrete_dtypes
@@ -841,7 +842,9 @@ def make_node(self, x, index):
         assert ind.ndim == 1
         assert ind.dtype in integer_dtypes
 
-        return Apply(self, [x, ind], [x.type()])
+        # The number of rows is set by the index, not by `x`.
+        out_shape = (ind.type.shape[0], x.type.shape[1])
+        return Apply(self, [x, ind], [x.type.clone(shape=out_shape)()])
 
     def perform(self, node, inp, outputs):
         (out,) = outputs
@@ -1113,7 +1116,14 @@ def make_node(self, x, index):
         if len(index) == 1:
             input_op += [generic_None, generic_None, generic_None]
 
-        return Apply(self, input_op, [x.type()])
+        # Both output dims are set by the slices, not by `x`, so derive each
+        # from its slice; a single index leaves the columns untouched.
+        out_shape = [
+            slice_static_length(ind, dim) for ind, dim in zip(index, x.type.shape)
+        ]
+        if len(index) == 1:
+            out_shape.append(x.type.shape[1])
+        return Apply(self, input_op, [x.type.clone(shape=tuple(out_shape))()])
 
     def perform(self, node, inputs, outputs):
         (x, start1, stop1, step1, start2, stop2, step2) = inputs

pytensor/tensor/blas/gemm.py

@@ -175,7 +175,9 @@ def make_node(self, *inputs):
         if not z.dtype.startswith("float") and not z.dtype.startswith("complex"):
             raise TypeError(Gemm.E_float, (z.dtype))
 
-        output = z.type()
+        # `z` is broadcast up to `dot(x, y)`, so the output shape is set by
+        # `x`/`y`, not by `z`'s (possibly broadcastable) static shape.
+        output = z.type.clone(shape=(x.type.shape[0], y.type.shape[1]))()
         return Apply(self, inputs, [output])
 
     def perform(self, node, inp, out):

pytensor/tensor/fourier.py

@@ -89,15 +89,20 @@ def make_node(self, a, n, axis):
                 raise TypeError(
                     "Length of the transformed axis must be a strictly positive scalar"
                 )
+        # The transformed axis is resized to `n`; the other axes keep `a`'s
+        # static sizes. Both are only known when `axis` and `n` are constant.
+        n_static = n.data.item() if isinstance(n, TensorConstant) else None
+        if isinstance(axis, TensorConstant):
+            axis_val = axis.data.item()
+            out_shape = tuple(
+                n_static if i == axis_val else s for i, s in enumerate(a.type.shape)
+            )
+        else:
+            out_shape = (None,) * a.ndim
         return Apply(
             self,
             [a, n, axis],
-            [
-                TensorType(
-                    "complex128",
-                    shape=tuple(1 if s == 1 else None for s in a.type.shape),
-                )()
-            ],
+            [TensorType("complex128", shape=out_shape)()],
         )
 
     def infer_shape(self, node, in_shapes):

pytensor/tensor/linalg/inverse.py

@@ -179,7 +179,9 @@ def __init__(self, ind=2):
 
     def make_node(self, a):
         a = as_tensor_variable(a)
-        out = a.type()
+        # `tensorinv` rotates the axes: the last `ndim - ind` dims come first.
+        out_shape = a.type.shape[self.ind :] + a.type.shape[: self.ind]
+        out = a.type.clone(shape=out_shape)()
         return Apply(self, [a], [out])
 
     def perform(self, node, inputs, outputs):

tests/sparse/test_basic.py

@@ -13,6 +13,7 @@
 from pytensor.gradient import GradientError
 from pytensor.graph.basic import Apply
 from pytensor.graph.op import Op
+from pytensor.graph.replace import vectorize_graph
 from pytensor.sparse.basic import (
     CSC,
     CSM,
@@ -1188,6 +1189,54 @@ def test_GetItemList_wrong_index(self):
         with pytest.raises(IndexError):
             f(A[0])
 
+    def test_GetItemList_static_shape(self):
+        # The number of output rows is set by the index, not by `x`, so the
+        # parent's row count must not leak into the output's static shape.
+        _a, A = sparse_random_inputs("csr", (4, 5))
+        x = as_sparse_variable(A[0])
+        assert x.type.shape == (4, 5)
+
+        y = sparse.get_item_list(x, lvector("index"))
+        assert y.type.shape == (None, 5)
+
+        index = pt.as_tensor_variable(np.array([0, 1, 2], dtype=np.int64))
+        y = sparse.get_item_list(x, index)
+        assert y.type.shape == (3, 5)
+
+    def test_GetItemList_vectorize(self):
+        # Regression test for gh-2274: leaking the parent's row count into the
+        # static output shape made `vectorize_graph` insert a `SpecifyShape`
+        # that failed at runtime when `len(index)` differed from that count.
+        _a, A = sparse_random_inputs("csr", (4, 5))
+        x = as_sparse_variable(A[0])
+        index = lvector("index")
+
+        coef = vector("coef")
+        eta = sparse.structured_dot(sparse.get_item_list(x, index), coef[:, None])
+
+        batch_coeff = matrix("batch_coeff")
+        batch_eta = vectorize_graph(eta, replace={coef: batch_coeff})
+        f = pytensor.function([batch_coeff, index], batch_eta)
+
+        sel = np.array([0, 3, 1], dtype=np.int64)
+        coeff_val = np.random.default_rng(0).random((8, 5))
+        out = np.asarray(f(coeff_val, sel))
+        assert out.shape == (8, 3, 1)
+
+        # Each batch b is `x[sel] @ coeff_val[b]`, broadcast over the batch axis.
+        expected = (A[0].toarray()[sel] @ coeff_val.T).T[:, :, None]
+        np.testing.assert_allclose(out, expected)
+
+    def test_GetItem2d_static_shape(self):
+        # Both output dims are set by the slices, not by `x`, so `x`'s shape
+        # must not leak into the output's static shape.
+        _a, A = sparse_random_inputs("csr", (5, 4))
+        x = as_sparse_variable(A[0])
+        assert x.type.shape == (5, 4)
+
+        assert x[0:2, :].type.shape == (2, 4)
+        assert x[1:5:2, 0:3].type.shape == (2, 3)
+
     def test_get_item_list_grad(self):
         op = sparse.basic.GetItemList()
 

tests/tensor/linalg/test_inverse.py

@@ -180,6 +180,13 @@ def test_infer_shape(self):
             TensorInv,
         )
 
+    def test_static_shape(self):
+        # `tensorinv` rotates the axes, so the input's shape must not leak
+        # unrotated into the output's static shape.
+        A = tensor4("A", shape=(4, 6, 8, 3))
+        assert tensorinv(A, ind=2).type.shape == (8, 3, 4, 6)
+        assert tensorinv(A, ind=1).type.shape == (6, 8, 3, 4)
+
     def test_eval(self):
         A = self.A
         Ai = tensorinv(A)

tests/tensor/test_blas.py

@@ -936,6 +936,16 @@ def test_gemm_broadcasting(inplace, linker):
         )
 
 
+def test_gemm_static_shape():
+    # `z` broadcasts up to `dot(x, y)`, so a broadcastable `z` must not leak
+    # its size-1 static shape into the output.
+    a, b = scalars("a", "b")
+    z = matrix("z", shape=(1, 1))
+    x = matrix("x", shape=(5, 4))
+    y = matrix("y", shape=(4, 3))
+    assert gemm_no_inplace(z, a, x, y, b).type.shape == (5, 3)
+
+
 def test_dot22():
     for dtype1 in ["float32", "float64", "complex64", "complex128"]:
         a = matrix(dtype=dtype1)

tests/tensor/test_fourier.py

@@ -3,7 +3,7 @@
 
 import pytensor
 from pytensor.tensor.fourier import Fourier, fft
-from pytensor.tensor.type import dmatrix, dvector, iscalar
+from pytensor.tensor.type import dmatrix, dvector, iscalar, tensor
 from tests import unittest_tools as utt
 
 
@@ -21,6 +21,15 @@ def test_perform(self):
         a = np.random.random((8, 6))
         assert np.allclose(f(a), np.fft.fft(a, 10, 0))
 
+    def test_static_shape(self):
+        # The transformed axis is resized to `n`; a size-1 input dim there must
+        # not leak into the output's static shape.
+        a = tensor("a", shape=(3, 1))
+        assert self.op(a, n=4, axis=1).type.shape == (3, 4)
+        # Unknown `n` or `axis` stays conservative.
+        assert self.op(a, n=iscalar(), axis=1).type.shape == (3, None)
+        assert self.op(a, n=4, axis=iscalar()).type.shape == (None, None)
+
     def test_infer_shape(self):
         a = dvector()
         self._compile_and_check(