BUG: convolve1d mode="same" indexing error

[tvwenger]

Describe the issue:

The new rewrite of convolve1d in #1318 and discussed in #1319 introduced a bug for mode="same". The documentation says:

        A string indicating the size of the output:
        - 'full': The output is the full discrete linear convolution of the inputs, with shape (..., N+M-1,).
        - 'valid': The output consists only of elements that do not rely on zero-padding, with shape (..., max(N, M) - min(N, M) + 1,).
        - 'same': The output is the same size as in1, centered with respect to the 'full' output.

Here it is the first axis of in2 that is compared for padding, rather than the intended last axis.

pytensor/pytensor/tensor/signal/conv.py

Lines 119 to 130 in 3af923b

	if mode == "same":
	# We implement "same" as "valid" with padded `in1`.
	in1_batch_shape = tuple(in1.shape)[:-1]
	zeros_left = in2.shape[0] // 2
	zeros_right = (in2.shape[0] - 1) // 2
	in1 = join(
	-1,
	zeros((*in1_batch_shape, zeros_left), dtype=in2.dtype),
	in1,
	zeros((*in1_batch_shape, zeros_right), dtype=in2.dtype),
	)
	mode = "valid"

Reproducable code example:

from pytensor.tensor import matrix, vector
from pytensor.tensor.signal.conv import convolve1d

def test_convolve1d_same():
    x = matrix("data")
    y = matrix("kernel")
    out = convolve1d(x, y, mode="same")

    rng = np.random.default_rng(38)
    x_test = rng.normal(size=(2, 8)).astype(x.dtype)
    y_test = rng.normal(size=(2, 8)).astype(x.dtype)

    res = out.eval({x: x_test, y: y_test})
    assert res.shape == (2, 8)

Error message:

PyTensor version information:

2.30.1

Context for the issue:

No response

[jessegrabowski]

Closed by #1337