def _test_gather_random_array(test_case, device):
input = np.random.randn(3, 4, 3, 5)
index = np.random.choice(np.arange(3), size=180, replace=True).reshape(
(3, 4, 3, 5))
np_out = np.take_along_axis(input, index, 1)
output = flow.gather(
flow.Tensor(input, device=flow.device(device)),
flow.Tensor(index, dtype=flow.int, device=flow.device(device)),
dim=1,
)
test_case.assertTrue(np.allclose(output.numpy(), np_out))
np_out2 = np.take_along_axis(input, index, 2)
output2 = flow.gather(
flow.Tensor(input, device=flow.device(device)),
flow.Tensor(index, dtype=flow.int, device=flow.device(device)),
dim=2,
)
test_case.assertTrue(np.allclose(output2.numpy(), np_out2))
np_out3 = np.take_along_axis(input, index, 3)
output3 = flow.gather(
flow.Tensor(input, device=flow.device(device)),
flow.Tensor(index, dtype=flow.int, device=flow.device(device)),
dim=3,
)
test_case.assertTrue(np.allclose(output3.numpy(), np_out3))
def test_gather(test_case):
input = np.array([[1, 2], [3, 4]])
index = np.array([[0, 0], [1, 0]])
np_out = np.take_along_axis(input, index, 0)
output = flow.gather(flow.Tensor(input),
flow.Tensor(index, dtype=flow.int),
dim=0)
test_case.assertTrue(np.array_equal(output.numpy(), np_out))
def test_gather_random_array(test_case):
input = np.random.randn(3, 4, 3, 5)
index = np.random.choice(np.arange(3), size=180, replace=True).reshape(
(3, 4, 3, 5)
)
np_out = gather_numpy(input, index, dim=1)
output = flow.gather(
flow.Tensor(input), flow.Tensor(index, dtype=flow.int), dim=1
)
test_case.assertTrue(np.allclose(output.numpy(), np_out))
np_out2 = gather_numpy(input, index, dim=2)
output2 = flow.gather(
flow.Tensor(input), flow.Tensor(index, dtype=flow.int), dim=2
)
test_case.assertTrue(np.allclose(output2.numpy(), np_out2))
np_out3 = gather_numpy(input, index, dim=3)
output3 = flow.gather(
flow.Tensor(input), flow.Tensor(index, dtype=flow.int), dim=3
)
test_case.assertTrue(np.allclose(output3.numpy(), np_out3))
def test_gather_backward(test_case):
input = np.array([[1, 2], [3, 4]])
index = np.array([[0, 0], [1, 0]])
np_out = np.take_along_axis(input, index, 0)
np_grad = _scatter_add_numpy(np.ones_like(np_out), 0, index,
input.shape)
of_input = flow.Tensor(input, requires_grad=True)
output = flow.gather(of_input,
flow.Tensor(index, dtype=flow.int),
dim=0)
out_sum = output.sum()
out_sum.backward()
test_case.assertTrue(np.array_equal(output.numpy(), np_out))
test_case.assertTrue(np.array_equal(of_input.grad.numpy(), np_grad))