Skip to content

Commit

Permalink
Allow int data-type for Embedding indices. (#6718)
Browse files Browse the repository at this point in the history
  • Loading branch information
ysiraichi authored Mar 12, 2024
1 parent 97acc14 commit 5a7bcac
Show file tree
Hide file tree
Showing 2 changed files with 17 additions and 1 deletion.
16 changes: 16 additions & 0 deletions test/test_operations.py
Original file line number Diff line number Diff line change
Expand Up @@ -1606,6 +1606,22 @@ def test_emb_bf16(self):
emb_out = emb(index)
assert emb_out.dtype == torch.bfloat16

def test_embedding_int_indices(self):
model = torch.nn.Embedding(1024, 10)

# 1 and 2-dimensional tensors.
# They have different execution paths.
for shape in ((5,), (2, 5)):

def test_on_device(device):
m = copy.deepcopy(model).to(device)
index = torch.ones(shape, dtype=torch.int, device=device)
return m(index)

out = test_on_device("cpu")
out_x = test_on_device(xm.xla_device())
self.assertEqual(out, out_x.cpu())

def test_transpose_1d(self):

def test_fn(t1):
Expand Down
2 changes: 1 addition & 1 deletion torch_xla/csrc/tensor_ops.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -245,7 +245,7 @@ XLATensorPtr EmbeddingDenseBackward(const XLATensorPtr& grad_output,
XLATensorPtr Embedding(const XLATensorPtr& weight,
const XLATensorPtr& indices) {
XLA_CHECK_EQ(weight->shape().get().rank(), 2);
XLA_CHECK_EQ(indices->dtype(), at::ScalarType::Long);
XLA_CHECK(indices->dtype() == at::kLong || indices->dtype() == at::kInt);

if (indices->shape().get().rank() == 1) {
return tensor_methods::index_select(weight, 0, indices);
Expand Down

0 comments on commit 5a7bcac

Please sign in to comment.