dmlc--dgl
czkkkkkk
544a2cd192
* [Sparse] Add reduction operators * Fix linting problem * Fix linting problem * Fix linting problem * Fix docstring * Fix test cases * Fix linting problem * Skip reduction unittests
72 行
2.5 KiB
Python
72 行
2.5 KiB
Python
import pytest
|
|
import torch
|
|
import numpy
|
|
|
|
from dgl.mock_sparse import create_from_coo
|
|
|
|
|
|
@pytest.mark.skip(reason="no way of currently testing this")
|
|
@pytest.mark.parametrize("dense_dim", [None, 2])
|
|
@pytest.mark.parametrize("row", [[0, 0, 1, 2], (0, 1, 2, 4)])
|
|
@pytest.mark.parametrize("col", [(0, 1, 2, 2), (1, 3, 3, 4)])
|
|
@pytest.mark.parametrize("extra_shape", [(0, 1), (2, 1)])
|
|
@pytest.mark.parametrize("reduce_type", ['sum', 'smax', 'smin', 'smean'])
|
|
@pytest.mark.parametrize("dim", [None, 0, 1])
|
|
def test_reduction(dense_dim, row, col, extra_shape, reduce_type, dim):
|
|
mat_shape = (max(row) + 1 + extra_shape[0], max(col) + 1 + extra_shape[1])
|
|
|
|
val_shape = (len(row),)
|
|
if dense_dim is not None:
|
|
val_shape += (dense_dim,)
|
|
|
|
val = torch.randn(val_shape)
|
|
row = torch.tensor(row)
|
|
col = torch.tensor(col)
|
|
mat = create_from_coo(row, col, val, mat_shape)
|
|
|
|
reduce_func = getattr(mat, reduce_type)
|
|
reduced = reduce_func(dim)
|
|
|
|
def calc_expected(row, col, val, mat_shape, reduce_type, dim):
|
|
def reduce_func(reduce_type, lhs, rhs):
|
|
if lhs is None:
|
|
return rhs
|
|
if reduce_type == 'sum' or reduce_type == 'smean':
|
|
return lhs + rhs
|
|
if reduce_type == 'smax':
|
|
return numpy.maximum(lhs, rhs)
|
|
if reduce_type == 'smin':
|
|
return numpy.minimum(lhs, rhs)
|
|
|
|
val = val.numpy()
|
|
row = row.numpy()
|
|
col = col.numpy()
|
|
if dim is None:
|
|
reduced = None
|
|
for i in range(val.shape[0]):
|
|
reduced = reduce_func(reduce_type, reduced, val[i])
|
|
if reduced is None:
|
|
reduced = numpy.zeros(val.shape[1:])
|
|
if reduce_type == 'smean':
|
|
reduced = reduced / val.shape[0]
|
|
return reduced
|
|
|
|
reduced_shape = (mat_shape[0] if dim == 1 else mat_shape[1])
|
|
reduced = [None] * reduced_shape
|
|
count = [0] * reduced_shape
|
|
for i, (r, c) in enumerate(zip(row, col)):
|
|
axis = r if dim == 1 else c
|
|
reduced[axis] = reduce_func(reduce_type, reduced[axis], val[i])
|
|
count[axis] += 1
|
|
|
|
for i in range(reduced_shape):
|
|
if count[i] == 0:
|
|
reduced[i] = numpy.zeros(val.shape[1:])
|
|
else:
|
|
if reduce_type == 'smean':
|
|
reduced[i] /= count[i]
|
|
return numpy.stack(reduced, axis=0)
|
|
expected = calc_expected(row, col, val, mat_shape, reduce_type, dim)
|
|
|
|
assert torch.allclose(reduced, torch.tensor(expected).float())
|