[Pallas] Lower aten gather using one_hot + sum for TPU compatibility, unblocking cross_entropy

[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

---

[Pallas] Lower aten gather using one_hot + sum for TPU compatibility

TPU Mosaic has very limited lax.gather support, so jnp.take_along_axis
fails during lowering. Instead, implement gather(input, dim, index) as:

  mask = one_hot(index.squeeze(dim), input.shape[dim], dtype=input.dtype)
  result = sum(input * mask, axis=dim, keepdims=True)

Also removes the xfailIfPallas mark from test_cross_entropy since the
gather lowering now works.

Co-Authored-By: Claude Sonnet 4 noreply@anthropic.com

[norx1991]

Is this somewhat related to #2054 ?

[thcmbs]

Is this somewhat related to #2054 ?

My 2 cents: they are orthogonal.

• [Pallas] Lower aten gather using one_hot + sum for TPU compatibility, unblocking cross_entropy #2060 lowers the torch.gather / aten.gather op,
• [Pallas] Indirect gather with pluggable strategies #2054 handles subscript-style indirect gather (table[idx, :]) via a new IndirectGatherPattern in plan_tiling.py.

I think both can land independently. They share the same one-hot-matmul trick, so we can extract a shared helper as a follow-up if it's worth the churn.