`OpFromGraph` falls back to object mode when output shape graph contains `Blockwise`

[jessegrabowski]

Description

When compiling a vectorized graph with an OpFromGraph into numba mode, shape inference fails if the shape of the output depends on blockwise operations. Example:

X = pt.dmatrix("X", shape=(None, None))
X_batched = pt.tensor("X", shape=(None, None, None))

Q, R = pt.linalg.qr(X)

results = OpFromGraph(
    inputs=[X],
    outputs=[Q],
)(X)

z_vec = vectorize_graph(results, {X: X_batched})
fn = pytensor.function(
    [X_batched],
    [z_vec],
    mode='NUMBA',
)

This falls back to object mode, because the introduce_explicit_core_shape_blockwise here forbids any blockwise. It should only be looking for loops -- specifically that the blockwise Op being rewritten is present in the shape graph.

In the above graph, here is the resulting shape graph:

Shape [id A] <Vector(int64, shape=(3,))>
 └─ Blockwise{OpFromGraph{inline=False}, (i00,i01)->(o00,o01)} [id B] <Tensor3(float64, shape=(?, ?, ?))>
    └─ X [id C] <Tensor3(float64, shape=(?, ?, ?))>

This might also be something unique to the QR Op? If so I'll edit the issue, but I still think the check in this rewrite is too strict.

[ricardoV94]

does qr have infer_shape?

[jessegrabowski]

No. I'll make an issue for that.

I think also hit this same bug with Join though, which does. The following graph drops to object mode:

    X = pt.dmatrix("X", shape=(10, None))
    X_batched = pt.tensor("X", shape=(None, 10, None))
    
    Q = pt.join(-1, X, pt.ones((10, 1)))
    
    Q_vec = pytensor.graph.replace.vectorize_graph(Q, {X: X_batched})
    fn = pytensor.function([X_batched],
                           Q_vec,
                           mode = 'NUMBA')

Graph of the shape of Q_vec:

Shape [id A]
 └─ Blockwise{Join, (),(i10,i11),(i20,i21)->(o00,o01)} [id B]
    ├─ ExpandDims{axis=0} [id C]
    │  └─ -1 [id D]
    ├─ X [id E]
    └─ ExpandDims{axis=0} [id F]
       └─ Alloc [id G]
          ├─ 1.0 [id H]
          ├─ 10 [id I]
          └─ 1 [id J]

[ricardoV94]

The infer shape is needed for sure, but no idea what's going on with join.

We can also tweak numba Blockwise to eval once in order to get the core shapes but if we can get it to just infer shapes that's cleaner