Minor cleanups to the gather expander; NFC

Minor cleanups to the gather expander; NFC

This change is NFC now, but it makes the code more general and this generality
will be used later on.  For instance

  ExpandFirstDimIntoNDims(transposed_gather_indices, {1, shape.dimensions(0)})

does not work if shape is a scalar shape (and this fine because today shape is
never scalar) but

  PrependDegenerateDims(transposed_gather_indices, 1)

works fine if transposed_gather_indices is scalar (and it will be, in a future
change).

PiperOrigin-RevId: 193283404