diff --git a/pytensor/sparse/basic.py b/pytensor/sparse/basic.py index 26a403729f..24424ebca6 100644 --- a/pytensor/sparse/basic.py +++ b/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 diff --git a/pytensor/tensor/blas/gemm.py b/pytensor/tensor/blas/gemm.py index 340dfadcaa..adfeedffd1 100644 --- a/pytensor/tensor/blas/gemm.py +++ b/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): diff --git a/pytensor/tensor/fourier.py b/pytensor/tensor/fourier.py index 260fc2dc09..1696e04d3b 100644 --- a/pytensor/tensor/fourier.py +++ b/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): diff --git a/pytensor/tensor/linalg/inverse.py b/pytensor/tensor/linalg/inverse.py index 6c205d6ae7..c5122473f9 100644 --- a/pytensor/tensor/linalg/inverse.py +++ b/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): diff --git a/tests/sparse/test_basic.py b/tests/sparse/test_basic.py index b3c663b5f0..98fb52b369 100644 --- a/tests/sparse/test_basic.py +++ b/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() diff --git a/tests/tensor/linalg/test_inverse.py b/tests/tensor/linalg/test_inverse.py index 471d3ddaeb..90a80a63c5 100644 --- a/tests/tensor/linalg/test_inverse.py +++ b/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) diff --git a/tests/tensor/test_blas.py b/tests/tensor/test_blas.py index b9ebc5939c..184c90f453 100644 --- a/tests/tensor/test_blas.py +++ b/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) diff --git a/tests/tensor/test_fourier.py b/tests/tensor/test_fourier.py index 088057ec0d..1eddd5b9d4 100644 --- a/tests/tensor/test_fourier.py +++ b/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(