Use torch.gather instead of generic int indexing for cross_entropy example

[AmesingFlank]

Stacked PRs:

• [Pallas] Switch gather to jnp.take_along_axis (for JAX issue filing) #2061
• [Pallas] Lower aten gather using one_hot + sum for TPU compatibility, unblocking cross_entropy #2060
• ->Use torch.gather instead of generic int indexing for cross_entropy example #2058

---

Use torch.gather instead of generic int indexing for cross_entropy example. The benefit of this is that the gather operation is applied on a tile, which is more friendly than an int-indexing expression on the entire tensor. Paired with #2060, this unblocks this cross_entropy example on TPUs

[norx1991]

Is there a way to know how this is affecting the GPU performance of this kernel? Or is the kernel code (e.g., triton) the same?

[jansel]

1. Let's benchmark this to make sure it doesn't regress GPU
2. Is it possible to make both versions work? We should at least add a test of the prior version so we can ensure it doesn't break on GPU in the future.