def test_getitem():
dims = [3]
sparse = SparseTensor(dims)
ref = np.zeros(dims)
sparse_slice = sparse[slice(2, 3, 2)]
ref_slice = ref[slice(2, 3, 2)]
assert sparse_slice.shape == ref_slice.shape
sparse_slice = sparse[slice(1, 3)][slice(1, 2, 2)]
ref_slice = ref[slice(1, 3)][slice(1, 2, 2)]
assert sparse_slice.shape == ref_slice.shape
dims = [2, 3]
sparse = SparseTensor(dims)
dense = DenseTensor(dims)
ref = np.zeros(dims)
sparse[1, -2] = 1
dense[1, -2] = 1
ref[1, -2] = 1
assert (sparse.to_nparray() == ref).all()
assert (dense.to_nparray() == ref).all()
dims = [5, 3, 1, 10]
sparse = SparseTensor(dims)
dense = DenseTensor(dims)
ref = np.zeros(dims)
assert sparse[:].shape == ref[:].shape
assert dense[:].shape == ref[:].shape
assert sparse[:, :].shape == ref[:, :].shape
assert sparse[:, :].shape == ref[:, :].shape
assert sparse[:, :, 0].shape == ref[:, :, 0].shape
assert dense[:, :, 0].shape == ref[:, :, 0].shape
assert sparse[:, :,
slice(1, None, 2)].shape == ref[:, :,
slice(1, None, 2)].shape
assert dense[:, :, slice(1, None, 2)].shape == ref[:, :,
slice(1, None, 2)].shape
sparse_slice = sparse[slice(2), :][slice(1, None, 2)]
dense_slice = dense[slice(2), :][slice(1, None, 2)]
ref_slice = ref[slice(2), :][slice(1, None, 2)]
assert sparse_slice.shape == ref_slice.shape
assert (sparse_slice == dense_slice).all()
for coord in [(3, 1, 0, 9), (3, -1, 0, 1)]:
sparse[coord] = 1
dense[coord] = 1
ref[coord] = 1
assert sparse[coord] == ref[coord]
assert dense[coord] == ref[coord]
assert (sparse.to_nparray() == ref).all()
assert (dense.to_nparray() == ref).all()
coord = (3, slice(1, 3))
assert (sparse[coord].to_nparray() == ref[coord]).all()
assert (dense[coord].to_nparray() == ref[coord]).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_eq():
dims = [3]
a = SparseTensor(dims)
b = SparseTensor(dims)
assert (a == b).all()
a = DenseTensor(dims)
b = DenseTensor(dims)
assert (a == b).all()
dims = [2, 1]
sparse = SparseTensor(dims)
dense = DenseTensor(dims)
ref = np.zeros(dims)
sparse[slice(1, None, 1)] = 3
dense[slice(1, None, 1)] = 3
ref[slice(1, None, 1)] = 3
assert (sparse == dense).all()
assert (sparse[0].to_nparray() == sparse.to_nparray()[0]).all()
assert (dense[0].to_nparray() == dense.to_nparray()[0]).all()
dims = [5, 7, 1, 12]
a = SparseTensor(dims)
b = SparseTensor(dims)
c = DenseTensor(dims)
d = DenseTensor(dims)
a[0, slice(1, -3, 2), :, slice(None, None, 4)] = 2
b[0, slice(1, -3, 2), :, slice(None, None, 4)] = 2
c[0, slice(1, -3, 2), :, slice(None, None, 4)] = 2
d[0, slice(1, -3, 2), :, slice(None, None, 4)] = 2
assert (a == b).all()
assert (c == d).all()
assert ((a == b) == (c == d)).all()
assert (a == c).all()
def test_expand_dims():
dims = [3, 1, 5, 2]
sparse = SparseTensor(dims)
dense = DenseTensor(dims)
ref = np.zeros(dims)
sparse[1, slice(None), slice(1, 5, 2), (0, 1)] = 1
dense[1, slice(None), slice(1, 5, 2), (0, 1)] = 1
ref[1, slice(None), slice(1, 5, 2), (0, 1)] = 1
assert (sparse.to_nparray() == ref).all()
assert (dense.to_nparray() == ref).all()
for axis in [3, 1, 0]:
sparse = sparse.expand_dims(axis)
dense = dense.expand_dims(axis)
ref = np.expand_dims(ref, axis)
assert (sparse.to_nparray() == ref).all()
assert (dense.to_nparray() == ref).all()
sparse = sparse.expand_dims(sparse.ndim)
dense = dense.expand_dims(dense.ndim)
ref = np.expand_dims(ref, ref.ndim)
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()
def test_setitem():
dims = [3, 4]
sparse = SparseTensor(dims)
dense = DenseTensor(dims)
ref = np.zeros(dims)
sparse[0, 1] = 1
dense[0, 1] = 1
ref[0, 1] = 1
assert (sparse == dense).all()
assert (sparse.to_nparray() == ref).all()
assert (dense.to_nparray() == ref).all()
sparse_slice = sparse[slice(None, None, 2)][slice(None, -1, 3)]
dense_slice = sparse[slice(None, None, 2)][slice(None, -1, 3)]
ref_slice = ref[slice(None, None, 2)][slice(None, -1, 3)]
assert sparse_slice.shape == ref_slice.shape
assert dense_slice.shape == ref_slice.shape
sparse[slice(None, None, 2)][slice(None, None, 3)] = 1
dense[slice(None, None, 2)][slice(None, None, 3)] = 1
ref[slice(None, None, 2)][slice(None, None, 3)] = 1
assert (sparse == dense).all()
assert (sparse.to_nparray() == ref).all()
assert (dense.to_nparray() == ref).all()
dims = [5, 7, 3, 10]
sparse = SparseTensor(dims)
dense = DenseTensor(dims)
ref = np.zeros(dims)
sparse[3, 0, 1, 9] = 1
dense[3, 0, 1, 9] = 1
ref[3, 0, 1, 9] = 1
assert (sparse == dense).all()
assert (sparse.to_nparray() == ref).all()
assert (dense.to_nparray() == ref).all()
zero_out = (2, 0, slice(None), slice(1, -3, 3))
sparse[zero_out] = 0
dense[zero_out] = 0
ref[zero_out] = 0
assert (sparse == dense).all()
assert (sparse.to_nparray() == ref).all()
assert (dense.to_nparray() == ref).all()
update = SparseTensor([3, 1])
update[2, 0] = 1
sparse[1] = update
dense[1] = update
ref[1] = update.to_nparray()
assert (sparse == dense).all()
assert (sparse.to_nparray() == ref).all()
assert (dense.to_nparray() == ref).all()
sparse[:, :, slice(1, None, 2)][slice(1, -3, 3)] = 3
dense[:, :, slice(1, None, 2)][slice(1, -3, 3)] = 3
ref[:, :, slice(1, None, 2)][slice(1, -3, 3)] = 3
assert (sparse == dense).all()
assert (sparse.to_nparray() == ref).all()
assert (dense.to_nparray() == ref).all()
sparse[0, slice(None, None, 2)][slice(None, None, 3)] = 4
dense[0, slice(None, None, 2)][slice(None, None, 3)] = 4
ref[0, slice(None, None, 2)][slice(None, None, 3)] = 4
assert (sparse == dense).all()
assert (sparse.to_nparray() == ref).all()
assert (dense.to_nparray() == ref).all()
sparse_slice = sparse[0]
dense_slice = dense[0]
ref_slice = ref[0]
assert (sparse == dense).all()
assert (sparse_slice.to_nparray() == ref_slice).all()
assert (dense_slice.to_nparray() == ref_slice).all()
sparse_slice = sparse[0, 0, :, :]
dense_slice = dense[0, 0, :, :]
ref_slice = ref[0, 0, :, :]
assert (sparse_slice == dense_slice).all()
assert (sparse_slice.to_nparray() == ref_slice).all()
assert (dense_slice.to_nparray() == ref_slice).all()
assert (sparse == dense).all()
assert (sparse.to_nparray() == ref).all()
assert (dense.to_nparray() == ref).all()