Infershape would incorrectly assign a shape of 1 on the last dimension,
if all dimension were preserved.
In particular, calling .flatten() on some vector would infer a shape of
(1,), which is incorrect.