This gives cleaner code and allows shape optimizations to happen on the result. PiperOrigin-RevId: 362242975
For now, the pass only reifies the required shape computations. Moving broadcasts will follow to allow for fusion across them. PiperOrigin-RevId: 362033715