Blockwise of an OpFromGraph core keeps the inner graph's unknown static shape

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))

Run the supplied reproduction and inspect vectorize_graph, Blockwise, and OpFromGraph handling of static shapes. Trace why the inner graph retains unknown dimensions instead of using the replaced input shape; done means the batched diag reports (?, 3, 3) and Shape_i of its final dimension folds to the constant 3.