import torch
import torchax
torchax.enable_globally()

T=1024  # Total number of tokens
D=2048  # hidden_size
L=128  # lora_rank
N=16  # num_loras

loras=torch.randn((N, L, D), dtype=torch.float32, device='jax')
idxs=torch.randint(0, N, (T,), dtype=torch.long, device='jax')
torch.nn.functional.one_hot(idxs, loras.shape[0])

>>> torch.nn.functional.one_hot(idxs, loras.shape[0])
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/mnt/disks/persist/github/xla/torchax/torchax/tensor.py", line 247, in __torch_function__
    return func(*args, **(kwargs or {}))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: one_hot is only applicable to index tensor of type LongTensor.

>>> torch.nn.functional.one_hot(idxs, loras.shape[0])
Tensor(<class 'jaxlib._jax.ArrayImpl'> [[0. 0. 0. ... 0. 0. 0.]
 [0. 0. 0. ... 0. 1. 0.]
 [0. 0. 0. ... 0. 0. 0.]
 ...
 [0. 0. 0. ... 0. 0. 0.]
 [0. 0. 0. ... 0. 0. 0.]
 [0. 0. 0. ... 0. 0. 1.]])