Add tests for uneven axis index groups in collectives.
* Fix a bug in psum and pmean with uneven axis index groups. We ignored raggedness when handling psum of a constant.
* Improve the errors for uneven axis index groups.
PiperOrigin-RevId: 919828788
