[Pallas] Introduce _make_group_metadata (#7107)
Summary:
_make_group_metadata is a helper function to assist gmm. Before we use the JAX version that cannot be stitched to our HLO. Now with this new torch version, it allows us to contain it in our HLO. However, we still need to lower two ops: pytorch.org/docs/stable/generated/torch.repeat_interleave.html and pytorch.org/docs/stable/generated/torch.histc.html. Yet we need the JAX version of the repeat to make the op shape static: jax.readthedocs.io/en/latest/_autosummary/jax.numpy.repeat.html.
Test Plan:
python test/test_gmm.py
