[Pallas] Support _histogram (#7115)
Summary:
This pull request implements a limited version of pytorch.org/docs/stable/generated/torch.histc.html to support _make_group_metadata. In the future, we can consider properly lower the op or just make this python version more completed.
Test Plan:
PJRT_DEVICE=TPU python test/test_gmm.py
