vectorize_graph over pt.diag returns a Blockwise{AllocDiag} whose output shape is (?, ?, ?) when the batched input has static core shape (3,), and Shape_i of that output never folds to a constant. AllocDiag is an OpFromGraph, and Blockwise reads the static shape from the inner graph, which was built before the input had a static shape, instead of recomputing it from the batched input. A plain Op core such as Cholesky recomputes and gets (?, 3, 3).
import pytensor.tensor as pt
from pytensor.graph import vectorize_graph
from pytensor.compile.mode import get_mode
from pytensor.graph import FunctionGraph
x = pt.vector("x") # static shape unknown when diag is built
d = pt.diag(x)
xb = pt.matrix("xb")
db = vectorize_graph(d, replace={x: pt.specify_shape(xb, (None, 3))})
print(db.type) # Tensor3(float64, shape=(?, ?, ?)), expected (?, 3, 3)
fg = FunctionGraph([xb], [db.shape[2]], clone=True)
get_mode("FAST_RUN").optimizer.rewrite(fg)
print(fg.outputs[0]) # Shape_i{2}.0, expected the constant 3
# Cholesky, a plain Op, recomputes from the batched input:
m = pt.matrix("m")
cb = vectorize_graph(pt.linalg.cholesky(m), replace={m: pt.specify_shape(pt.tensor3("mb"), (None, 3, 3))})
print(cb.type) # Tensor3(float64, shape=(?, 3, 3))
vectorize_graphoverpt.diagreturns aBlockwise{AllocDiag}whose output shape is(?, ?, ?)when the batched input has static core shape(3,), andShape_iof that output never folds to a constant.AllocDiagis anOpFromGraph, andBlockwisereads the static shape from the inner graph, which was built before the input had a static shape, instead of recomputing it from the batched input. A plainOpcore such asCholeskyrecomputes and gets(?, 3, 3).