Revert line from #4585 to get mask, data shapes to match in .flat