`Blockwise{Convolve1d}` never infers static output shape

[ErikRingen]

Describe the issue:

pt.signal.convolve1d always produces (None,) output shape, even when both input shapes and the mode are fully known at graph construction time.

Reproducable code example:

import pytensor.tensor as pt

x = pt.vector('x', shape=(61,))
k = pt.vector('k', shape=(10,))

conv = pt.signal.convolve1d(x, k, mode='full')
print(conv.type.shape)  # (None,)  — expected (70,)


bx = pt.matrix('bx', shape=(4, 61))
bk = pt.matrix('bk', shape=(4, 10))
bconv = pt.signal.convolve1d(bx, bk, mode='full')
print(bconv.type.shape)  # (4, None)  — expected (4, 70)

Error message:

PyTensor version information:

• PyTensor version: 2.38.2 (also confirmed on main at c31183051)
• Python version: 3.12
• Operating system: macOS (ARM64)
• How did you install PyTensor: pixi (conda-forge)

[ricardoV94]

It can be achieved, because there is a setting for Blockwise to build the core_node (from which it extracts shape info), using literal scalars, instead of dummies. However it shouldn't be strictly needed. We do that at dispatch time for JAX Blockwise already here:

pytensor/pytensor/link/jax/dispatch/blockwise.py

Line 11 in d7541cd

node.inputs, propagate_unbatched_core_inputs=True

I assume the issue is not necessarily Blockwise/Convolve1d but something else in the graph that looks like a dynamic shape for JAX but is not if we constant-folded more aggressively. Could you investigate what that might be, or provide a MWE that now fails to compile to JAX so we can investigate.

For reference, this compiles just fine to Jax:

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

bx = pt.matrix('bx', shape=(4, 61))
bk = pt.matrix('bk', shape=(4, 10))
out = pt.signal.convolve1d(bx, bk, mode='full')
grad_bx, grad_bx = pt.grad(out.sum(), wrt=[bx, bx])

rng = np.random.default_rng()
bx_test = rng.normal(size=bx.type.shape)
bk_test = rng.normal(size=bk.type.shape)

fn = pytensor.function([bx, bk], [out, grad_bx, grad_bx], mode="JAX")
fn(bx_test, bk_test)  # fine

In general PyTensor should work even if static shapes can't be inferred, and it's a code-complexity graph definition trade-off to compute them in advance. We can consider the special work here, but first would be good to find out what is the real issue that shows up in workflow.

[ricardoV94]

I'm removing the bug label, static shape is not a requirement, it's a nice to have. The real bug is something related to the way we compile JAX and it not liking this, but we need a MWE