def test_with_negative_1(self):
devices = [torch.device('cpu')]
if torch.cuda.is_available():
devices.append(torch.device('cuda', 0))
for device in devices:
src = torch.tensor([0, 1, 2, 3],
dtype=torch.float32,
requires_grad=True,
device=device)
indexes = k2.RaggedInt(
'[ [1 2 -1] [0 3] [-1] [0 2 3 1 3] [] ]').to(device)
ans = k2.index_and_sum(src, indexes)
expected = torch.tensor([1 + 2, 0 + 3, 0, 0 + 2 + 3 + 1 + 3,
0]).to(src)
assert torch.all(torch.eq(ans, expected))
# now for autograd
scale = torch.tensor([10, 20, 30, 40, 50]).to(device)
(ans * scale).sum().backward()
expected_grad = torch.empty_like(src.grad)
expected_grad[0] = scale[1] + scale[3]
expected_grad[1] = scale[0] + scale[3]
expected_grad[2] = scale[0] + scale[3]
expected_grad[3] = scale[1] + scale[3] * 2
assert torch.all(torch.eq(src.grad, expected_grad))
def test_without_negative_1(self):
for device in self.devices:
src = torch.tensor([0, 1, 2, 3],
dtype=torch.float32,
requires_grad=True,
device=device)
indexes = k2.RaggedInt('[ [1 2] [0 3] [0 2 3 1 3] [] ]').to(device)
ans = k2.index_and_sum(src, indexes)
expected = torch.tensor([1 + 2, 0 + 3, 0 + 2 + 3 + 1 + 3,
0]).to(src)
assert torch.all(torch.eq(ans, expected))
# now for autograd
scale = torch.tensor([10, 20, 30, 40]).to(device)
(ans * scale).sum().backward()
expected_grad = torch.empty_like(src.grad)
expected_grad[0] = scale[1] + scale[2]
expected_grad[1] = scale[0] + scale[2]
expected_grad[2] = scale[0] + scale[2]
expected_grad[3] = scale[1] + scale[2] * 2
assert torch.allclose(src.grad, expected_grad)