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.
Reverts 8ff771c767df9b213dd81b6642fa0d1fff0db885
PiperOrigin-RevId: 919928349
