embeddingbad error

[ailzhang]

reported by @taylanbil

import torch
import torch.nn as nn
import torch_xla.core.xla_model as xm


device = xm.xla_device()
d = nn.EmbeddingBag(10, 10, mode="sum", sparse=False).to(device)
inp = torch.LongTensor([1, 5, 9]).to(device)
x = d(inp, offsets=torch.LongTensor([0]).to(device))
loss = x.sum()
loss.backward()

[gmagogsfm]

Here is a repro:
https://gist.github.com/gmagogsfm/00d18a2972cafda3bd7d97fc6ece3354

Ailing, could you confirm if this is the issue you are seeing?

[ailzhang]

@gmagogsfm yup that’s it!

[gmagogsfm]

Discussed offline with ailzhang

Root cause is EmbeddingBag's kernel uses weight.requires_grad() to check wether outputs like bag_size will be computed

But when converting aten::Tensor to XLATensor, requires_grad bit is lost because XLATensor doesn't capture this information at all. This causes bag_size to always to undef tensor, ultimately resulting in error during backward prop.

Solution we plan to implement is splitting embedding_bag op into two different op:

1. embedding_bag_forward_only, that doesn't compute outputs required by grad computation
2. embedding_bag that computes all output

They should still be covered under same API nn.EmbeddingBag, so usability is not hurt.