[Mosaic TPU] Add a missing check in load->reshape optimization
Previously we assumed but didn't actually check that whole dimensions are flattened
into the minor dim. It should not be too difficult to generalize this, but for now it's
better to at least have a check. Without the modification the newly added test case crashes
due to a CHECK fail.
PiperOrigin-RevId: 828017228
