def test_transpose():
dims = [3, 2]
sparse = SparseTensor(dims)
dense = DenseTensor(dims)
ref = np.zeros(dims)
sparse[slice(2), 1] = 1
dense[slice(2), 1] = 1
ref[slice(2), 1] = 1
assert (sparse.transpose().to_nparray() == ref.transpose()).all()
assert (sparse.transpose([0, 1]).to_nparray() == ref).all()
assert (sparse.transpose([1,
0]).to_nparray() == np.transpose(ref,
[1, 0])).all()
assert (dense.transpose().to_nparray() == ref.transpose()).all()
assert (dense.transpose([0, 1]).to_nparray() == ref).all()
assert (dense.transpose([1,
0]).to_nparray() == np.transpose(ref,
[1, 0])).all()
dims = [5, 1, 8, 3]
sparse = SparseTensor(dims)
dense = DenseTensor(dims)
ref = np.zeros(dims)
sparse[:, :, slice(None, None, 2)] = 1
dense[:, :, slice(None, None, 2)] = 1
ref[:, :, slice(None, None, 2)] = 1
assert (sparse.transpose().to_nparray() == ref.transpose()).all()
assert (sparse.transpose([0, 2, 1, 3]).to_nparray() == np.transpose(
ref, [0, 2, 1, 3])).all()
assert (sparse.transpose([3, 1, 2, 0]).to_nparray() == np.transpose(
ref, [3, 1, 2, 0])).all()
assert (dense.transpose().to_nparray() == ref.transpose()).all()
assert (dense.transpose([0, 2, 1, 3]).to_nparray() == np.transpose(
ref, [0, 2, 1, 3])).all()
assert (dense.transpose([3, 1, 2, 0]).to_nparray() == np.transpose(
ref, [3, 1, 2, 0])).all()
permutation = [3, 1, 0, 2] # shape = (3, 1, 5, 8)
sparse_slice = sparse.transpose(permutation)[:, slice(None, 4), 0]
dense_slice = dense.transpose(permutation)[:, slice(None, 4), 0]
ref_slice = ref.transpose(permutation)[:, slice(None, 4), 0]
assert sparse_slice.shape == ref_slice.shape
assert dense_slice.shape == ref_slice.shape
assert (sparse_slice.to_nparray() == ref_slice).all()
assert (dense_slice.to_nparray() == ref_slice).all()
def test_sum():
dims = [2, 3, 4, 5]
sparse = SparseTensor(dims)
dense = DenseTensor(dims)
ref = np.zeros(dims)
sparse[:, slice(None, None, 2)] = 1
dense[:, slice(None, None, 2)] = 1
ref[:, slice(None, None, 2)] = 1
assert (sparse.sum(0).to_nparray() == ref.sum(0)).all()
assert (dense.sum(0).to_nparray() == ref.sum(0)).all()
assert (sparse.sum(2).to_nparray() == ref.sum(2)).all()
assert (dense.sum(2).to_nparray() == ref.sum(2)).all()
sparse_transpose = sparse.transpose()[:, 0].sum(2).expand_dims(1)
dense_transpose = dense.transpose()[:, 0].sum(2).expand_dims(1)
ref_transpose = np.expand_dims(np.transpose(ref)[:, 0].sum(2), 1)
assert (sparse_transpose.to_nparray() == ref_transpose).all()
assert (dense_transpose.to_nparray() == ref_transpose).all()
for axis in range(2):
sparse = sparse.sum(0)
dense = dense.sum(0)
ref = ref.sum(0)
assert (sparse.to_nparray() == ref).all()
assert (dense.to_nparray() == ref).all()
dims = [3, 5, 2, 4]
sparse = SparseTensor(dims)
dense = DenseTensor(dims)
ref = np.zeros(dims)
assert (sparse.sum(2).to_nparray() == np.sum(ref, 2)).all()
assert (dense.sum(2).to_nparray() == np.sum(ref, 2)).all()
sparse[:, :, 0, slice(None, None, 3)] = 2
dense[:, :, 0, slice(None, None, 3)] = 2
ref[:, :, 0, slice(None, None, 3)] = 2
for axis in range(4):
assert (sparse.sum(axis).to_nparray() == np.sum(ref, axis)).all()
assert (dense.sum(axis).to_nparray() == np.sum(ref, axis)).all()
for axis in range(3):
sparse = sparse.sum(0)
dense = dense.sum(0)
ref = ref.sum(0)
assert (sparse.to_nparray() == ref).all()
assert (dense.to_nparray() == ref).all()
def test_product():
dims = [3, 5, 2, 4]
sparse = SparseTensor(dims)
dense = DenseTensor(dims)
ref = np.zeros(dims)
assert (sparse.product(2).to_nparray() == np.product(ref, 2)).all()
assert (dense.product(2).to_nparray() == np.product(ref, 2)).all()
sparse[:, :, 0, slice(None, None, 3)] = 2
dense[:, :, 0, slice(None, None, 3)] = 2
ref[:, :, 0, slice(None, None, 3)] = 2
for axis in range(4):
assert (sparse.product(axis).to_nparray() == np.product(ref,
axis)).all()
assert (dense.product(axis).to_nparray() == np.product(ref,
axis)).all()
sparse = sparse.transpose()[:, 0].product(1).expand_dims(1)
dense = dense.transpose()[:, 0].product(1).expand_dims(1)
ref = np.expand_dims(np.product(np.transpose(ref)[:, 0], 1), 1)
assert (sparse.to_nparray() == ref).all()
assert (dense.to_nparray() == ref).all()