pt.flatten drops static shape on a fully-known input #2366

flatten builds its reshape target as (*x.shape[:ndim-1], -1), and unpacking a slice of the shape vector gives entries that don't constant-fold, so the output type loses shapes an equivalent reshape keeps.

import pytensor.tensor as pt

x = pt.tensor("x", shape=(32, 24, 24, 16))

print(pt.flatten(x, ndim=2).type.shape)        # (None, None)
print(x.reshape((x.shape[0], -1)).type.shape)  # (32, 9216)

Potential fix (requires testing):

dims = (*(_x.shape[i] for i in range(ndim - 1)), -1)

Indexing per axis yields Shape_i, which folds where the extent is known; a slice of the shape vector yields a Subtensor that doesn't. This recovers (32, 9216) but not the partly-dynamic case — with an unknown batch, Reshape still can't resolve the -1 even when every other extent is known, so it stays (None, None) where (None, 9216) is derivable.