def test_reduce_sum(self):
values, row_index, col_index = self._prepare_tables()
cell_index = ProductIndexMap(row_index, col_index)
row_sum, _ = reduce_sum(values, row_index)
col_sum, _ = reduce_sum(values, col_index)
cell_sum, _ = reduce_sum(values, cell_index)
# We use np.testing.assert_allclose rather than Tensorflow's assertAllClose
np.testing.assert_allclose(row_sum.numpy(), [[6.0, 3.0, 8.0], [6.0, 3.0, 8.0]])
np.testing.assert_allclose(col_sum.numpy(), [[9.0, 8.0, 0.0], [4.0, 5.0, 8.0]])
np.testing.assert_allclose(
cell_sum.numpy(),
[[3.0, 3.0, 0.0, 2.0, 1.0, 0.0, 4.0, 4.0, 0.0], [1.0, 2.0, 3.0, 2.0, 0.0, 1.0, 1.0, 3.0, 4.0]],
)
def test_reduce_sum_vectorized(self):
values = torch.as_tensor([[1.0, 2.0, 3.0], [2.0, 3.0, 4.0], [3.0, 4.0, 5.0]])
index = IndexMap(indices=torch.as_tensor([0, 0, 1]), num_segments=2, batch_dims=0)
sums, new_index = reduce_sum(values, index)
# We use np.testing.assert_allclose rather than Tensorflow's assertAllClose
np.testing.assert_allclose(sums.numpy(), [[3.0, 5.0, 7.0], [3.0, 4.0, 5.0]])
# We use np.testing.assert_array_equal rather than Tensorflow's assertAllEqual
np.testing.assert_array_equal(new_index.indices.numpy(), [0, 1])
np.testing.assert_array_equal(new_index.num_segments.numpy(), 2)
np.testing.assert_array_equal(new_index.batch_dims, 0)
def test_gather(self):
values, row_index, col_index = self._prepare_tables()
cell_index = ProductIndexMap(row_index, col_index)
# Compute sums and then gather. The result should have the same shape as
# the original table and each element should contain the sum the values in
# its cell.
sums, _ = reduce_sum(values, cell_index)
cell_sum = gather(sums, cell_index)
assert cell_sum.size() == values.size()
# We use np.testing.assert_array_equal rather than Tensorflow's assertAllEqual
np.testing.assert_allclose(
cell_sum.numpy(),
[[[3.0, 3.0, 3.0], [2.0, 2.0, 1.0], [4.0, 4.0, 4.0]], [[1.0, 2.0, 3.0], [2.0, 0.0, 1.0], [1.0, 3.0, 4.0]]],
)
