Google Git
Sign in
android / platform / external / pytorch / refs/heads/sdk-release / . / test / test_sparse.py
blob: 6b35b78200714591c576c4e5e0660c0a5e7fe908 [file] [log] [blame] [edit]
# Owner(s): ["module: sparse"]
import torch
import itertools
import functools
import operator
import random
import unittest
from torch.testing import make_tensor
from torch.testing._internal.common_utils import TestCase, run_tests, skipIfRocm, do_test_dtypes, \
load_tests, TEST_NUMPY, TEST_SCIPY, IS_WINDOWS, gradcheck, coalescedonoff, \
DeterministicGuard, first_sample, TEST_WITH_CROSSREF, TEST_WITH_ROCM, skipIfTorchDynamo, \
parametrize, subtest, is_coalesced_indices, suppress_warnings, instantiate_parametrized_tests, \
skipIfCrossRef
from torch.testing._internal.common_cuda import TEST_CUDA
from numbers import Number
from typing import Dict, Any
from packaging import version
from torch.testing._internal.common_cuda import \
(SM53OrLater, SM80OrLater, TEST_MULTIGPU)
from torch.testing._internal.common_device_type import \
(instantiate_device_type_tests, ops, dtypes, dtypesIfCUDA, onlyCPU, onlyCUDA, precisionOverride,
deviceCountAtLeast, OpDTypes, onlyNativeDeviceTypes)
from torch.testing._internal.common_methods_invocations import \
(op_db, reduction_ops, sparse_unary_ufuncs, sparse_masked_reduction_ops, binary_ufuncs)
from torch.testing._internal.common_dtype import (
all_types, all_types_and_complex, all_types_and_complex_and, floating_and_complex_types,
floating_and_complex_types_and, integral_types, floating_types_and,
)
from torch.testing._internal.opinfo.definitions.sparse import validate_sample_input_sparse
from torch.testing._internal.opinfo.refs import (
ElementwiseBinaryPythonRefInfo,
ReductionPythonRefInfo
)
def _op_supports_any_sparse(op):
return (op.supports_sparse
or op.supports_sparse_csr
or op.supports_sparse_csc
or op.supports_sparse_bsr
or op.supports_sparse_bsc)
reduction_ops_with_sparse_support = [
op for op in reduction_ops if 'masked.' not in op.name and
_op_supports_any_sparse(op) and not isinstance(op, ReductionPythonRefInfo)]
binary_ufuncs_with_sparse_support = [
op for op in binary_ufuncs if _op_supports_any_sparse(op) and
not isinstance(op, ElementwiseBinaryPythonRefInfo)]
like_fns_with_sparse_support = [op for op in op_db if _op_supports_any_sparse(op) and '_like' in op.name]
if TEST_SCIPY:
import scipy.sparse
# load_tests from torch.testing._internal.common_utils is used to automatically filter tests for
# sharding on sandcastle. This line silences flake warnings
load_tests = load_tests
# batched grad doesn't support sparse
gradcheck = functools.partial(gradcheck, check_batched_grad=False)
CUSPARSE_SPMM_COMPLEX128_SUPPORTED = (
IS_WINDOWS and torch.version.cuda and version.parse(torch.version.cuda) > version.parse("11.2")
) or (not IS_WINDOWS and not TEST_WITH_ROCM)
HIPSPARSE_SPMM_COMPLEX128_SUPPORTED = torch.version.hip and version.parse(torch.version.hip.split("-")[0]) >= version.parse("6.0")
def all_sparse_layouts(test_name='layout', include_strided=False):
return parametrize(test_name, [
subtest(torch.strided, name='Strided'),
subtest(torch.sparse_coo, name='SparseCOO'),
subtest(torch.sparse_csr, name='SparseCSR'),
subtest(torch.sparse_csc, name='SparseCSC'),
subtest(torch.sparse_bsr, name='SparseBSR'),
subtest(torch.sparse_bsc, name='SparseBSC'),
][(0 if include_strided else 1):])
def gradcheck_semantics(test_name='gradcheck'):
gradcheck_sparse = functools.partial(gradcheck, masked=False)
gradcheck_masked = functools.partial(gradcheck, masked=True)
gradcheck_sparse.masked = False
gradcheck_masked.masked = True
return parametrize(test_name, [
subtest(gradcheck_sparse, name='sparse'),
subtest(gradcheck_masked, name='masked')])
class CrossRefSparseFakeMode(torch._subclasses.CrossRefFakeMode):
def __init__(self) -> None:
super().__init__(
self.ignore_op, check_strides=False,
check_aliasing=False,
) # TODO: enable stride/alias checking
# empty_like excluded for now due to sparse complex
# aten._to_dense.default this one is getting called with csc
@staticmethod
def ignore_op(func):
return func in (
torch.ops.aten.empty_like.default,
torch.ops.aten.set_.source_Storage_storage_offset,
torch.ops.aten.sspaddmm.out,
torch.ops.aten._spdiags.default,
torch.ops.aten._to_dense.default,
torch.ops.aten.indices.default,
torch.ops.aten._indices.default,
torch.ops.aten.values.default,
torch.ops.aten._values.default,
)
class TestSparseLegacyAndDeprecation(TestCase):
@skipIfTorchDynamo("TorchDynamo fails with unknown reason")
def test_legacy_warnings(self):
def f1():
"torch.sparse.SparseTensor() is deprecated."\
" Please use torch.sparse_coo_tensor((0,), dtype=)"
x_ref = torch.sparse_coo_tensor((0,), dtype=torch.float64)
x = torch.sparse.DoubleTensor()
self.assertEqual(x, x_ref)
def f2():
"torch.sparse.SparseTensor(cdata=x._cdata) is deprecated."\
" Please use torch.sparse_coo_tensor(x._indices(), x._values(), x.shape)"
x_ref = torch.tensor([[1, 2], [3, 4]], dtype=torch.float64).to_sparse()
x = torch.sparse.DoubleTensor(cdata=x_ref._cdata)
y = torch.sparse_coo_tensor(x._indices(), x._values(), x.shape)
self.assertEqual(x, x_ref)
self.assertEqual(y, x_ref)
def f3():
"torch.sparse.SparseTensor(indices, values, *, device=) is deprecated."\
" Please use torch.sparse_coo_tensor(indices, values, dtype=, device=)"
x_ref = torch.sparse_coo_tensor([[0, 0, 1, 1], [0, 1, 0, 1]], [1, 2, 3, 4], dtype=torch.float64)
x = torch.sparse.DoubleTensor(torch.tensor([[0, 0, 1, 1], [0, 1, 0, 1]]),
torch.tensor([1, 2, 3, 4], dtype=torch.float64))
self.assertEqual(x, x_ref)
def f4():
"torch.sparse.SparseTensor(indices, values, shape, *, device=) is deprecated."\
" Please use torch.sparse_coo_tensor(indices, values, shape, dtype=, device=)"
x_ref = torch.sparse_coo_tensor([[0, 0, 1, 1], [0, 1, 0, 1]], [1, 2, 3, 4], (2, 3), dtype=torch.float64)
x = torch.sparse.DoubleTensor(torch.tensor([[0, 0, 1, 1], [0, 1, 0, 1]]),
torch.tensor([1, 2, 3, 4], dtype=torch.float64), (2, 3))
self.assertEqual(x, x_ref)
def f5():
"torch.sparse.SparseTensor(shape, *, device=) is deprecated."\
" Please use torch.sparse_coo_tensor(shape, dtype=, device=)"
x_ref = torch.sparse_coo_tensor((2, 3), dtype=torch.float64)
x = torch.sparse.DoubleTensor(2, 3)
self.assertEqual(x, x_ref)
for test_f in [f1, f2, f3, f4, f5]:
with self.assertWarns(UserWarning, msg=test_f.__doc__) as cm:
test_f()
test_f()
# Check warn-once:
self.assertEqual(len(cm.warnings), 1)
class TestSparseBase(TestCase):
def run(self, result=None):
if TEST_WITH_CROSSREF:
with CrossRefSparseFakeMode():
return super().run(result)
else:
return super().run(result)
class TestSparse(TestSparseBase):
def setUp(self):
TestCase.setUp(self)
self.index_tensor = lambda *args, **kwargs: torch.tensor(*args, **kwargs, dtype=torch.int64)
def sparse_empty_factory(*args, **kwargs):
kwargs['layout'] = kwargs.get('layout', torch.sparse_coo)
return torch.empty(*args, **kwargs)
self.sparse_empty = sparse_empty_factory
def sparse_tensor_factory(*args, **kwargs):
return torch.sparse_coo_tensor(*args, **kwargs)
self.sparse_tensor = sparse_tensor_factory
def _gen_sparse(self, sparse_dim, nnz, with_size, dtype, device, coalesced):
if isinstance(with_size, Number):
with_size = [with_size] * sparse_dim
x, i, v = self.genSparseTensor(with_size, sparse_dim, nnz, not coalesced, dtype=dtype, device=device)
if not coalesced:
self.assert_uncoalesced(x)
return x, i, v
def assert_uncoalesced(self, x):
"""
Test if a CPU tensor is uncoalesced. This is used to ensure
correctness of the uncoalesced tensor generation algorithm.
"""
assert not x.is_coalesced()
existing_indices = set()
indices = x._indices()
for i in range(x._nnz()):
index = str(indices[:, i])
if index in existing_indices:
return True
else:
existing_indices.add(index)
def randn(self, *args, **kwargs):
"""
Variant of torch.randn that also works in the TEST_CUDA case.
"""
# TODO: Put this in torch.cuda.randn
return torch.empty(*args, **kwargs).normal_()
@dtypes(torch.double)
def test_print_coalesced(self, device, dtype):
self._test_print(device, dtype, True)
@dtypes(torch.double)
def test_print_uncoalesced(self, device, dtype):
self._test_print(device, dtype, False)
def _test_print(self, device, dtype, coalesced):
shape_sparse_dim_nnz = [
((), 0, 2),
((0,), 0, 10),
((2,), 0, 3),
((100, 3), 1, 3),
((100, 20, 3), 2, 0),
((10, 0, 3), 0, 3),
((10, 0, 3), 0, 0),
]
printed = []
for shape, sparse_dim, nnz in shape_sparse_dim_nnz:
indices_shape = torch.Size((sparse_dim, nnz))
values_shape = torch.Size((nnz,) + shape[sparse_dim:])
printed.append(f"# shape: {torch.Size(shape)}")
printed.append(f"# nnz: {nnz}")
printed.append(f"# sparse_dim: {sparse_dim}")
printed.append(f"# indices shape: {indices_shape}")
printed.append(f"# values shape: {values_shape}")
indices = torch.arange(indices_shape.numel(), dtype=self.index_tensor(0).dtype,
device=device).view(indices_shape)
for d in range(sparse_dim):
indices[d].clamp_(max=(shape[d] - 1)) # make it valid index
if not coalesced and indices.numel() > 0:
indices[:, -1] = indices[:, 0] # make it uncoalesced
values_numel = values_shape.numel()
values = torch.arange(values_numel, dtype=dtype,
device=device).view(values_shape).div_(values_numel / 2.)
sp_tensor = self.sparse_tensor(indices, values, shape, dtype=dtype, device=device)
dtypes = [torch.int32]
if values.dtype == torch.double:
dtypes.append(torch.float)
else:
dtypes.append(torch.double)
for dtype in dtypes:
printed.append(f"########## {dtype} ##########")
x = sp_tensor.detach().to(dtype)
printed.append("# sparse tensor")
printed.append(str(x))
if x.dtype.is_floating_point:
printed.append("# after requires_grad_")
printed.append(str(x.requires_grad_()))
printed.append("# after addition")
printed.append(str(x + x))
printed.append("# _indices")
printed.append(str(x._indices()))
printed.append("# _values")
printed.append(str(x._values()))
printed.append('')
self.assertExpected('\n'.join(printed))
@coalescedonoff
@dtypes(torch.double, torch.cdouble)
def test_basic(self, device, dtype, coalesced):
def test_shape(sparse_dims, nnz, with_size):
if isinstance(with_size, Number):
with_size = [with_size] * sparse_dims
x, i, v = self._gen_sparse(sparse_dims, nnz, with_size, dtype, device, coalesced)
self.assertEqual(i, x._indices())
self.assertEqual(v, x._values())
self.assertEqual(x.ndimension(), len(with_size))
self.assertEqual(x.coalesce()._nnz(), nnz if x.is_coalesced() else nnz // 2)
self.assertEqual(list(x.size()), with_size)
# Test .indices() and .values()
if not coalesced:
with self.assertRaisesRegex(RuntimeError, "Cannot get indices on an uncoalesced tensor"):
x.indices()
with self.assertRaisesRegex(RuntimeError, "Cannot get values on an uncoalesced tensor"):
x.values()
else:
self.assertEqual(x.indices(), x._indices())
self.assertEqual(x.values(), x._values())
test_shape(3, 10, 100)
test_shape(3, 10, [100, 100, 100])
test_shape(3, 10, [100, 100, 100, 5, 5, 5, 0])
test_shape(3, 0, [0, 0, 100, 5, 5, 5, 0])
# Make sure that coalesce handles duplicate indices correctly
i = self.index_tensor([[9, 0, 0, 0, 8, 1, 1, 1, 2, 7, 2, 2, 3, 4, 6, 9]], device=device)
v = torch.tensor([[idx**2, idx] for idx in range(i.size(1))], dtype=dtype, device=device)
x = self.sparse_tensor(i, v, torch.Size([10, 2]), dtype=dtype, device=device)
self.assertEqual(x.coalesce()._nnz(), 9)
@coalescedonoff
@dtypes(torch.double, torch.cdouble, torch.bfloat16)
@precisionOverride({torch.bfloat16: 1e-2})
@skipIfTorchDynamo("https://github.com/pytorch/torchdynamo/issues/1991")
def test_coalesce(self, device, dtype, coalesced):
def _test_coalesce(t):
tc = t.coalesce()
self.assertEqual(tc.to_dense(), t.to_dense())
self.assertTrue(tc.is_coalesced())
# Our code below doesn't work when nnz is 0, because
# then it's a 0D tensor, not a 2D tensor.
if t._nnz() == 0:
self.assertEqual(t._indices(), tc._indices())
self.assertEqual(t._values(), tc._values())
return tc
value_map: Dict[Any, Any] = {}
for idx, val in zip(t._indices().t(), t._values()):
idx_tup = tuple(idx.tolist())
if idx_tup in value_map:
value_map[idx_tup] += val
else:
value_map[idx_tup] = val.clone() if isinstance(val, torch.Tensor) else val
new_indices = sorted(value_map.keys())
_new_values = [value_map[idx] for idx in new_indices]
if t._values().ndimension() < 2:
new_values = t._values().new(_new_values)
else:
new_values = torch.stack(_new_values)
new_indices = t._indices().new(new_indices).t()
tg = t.new(new_indices, new_values, t.size())
self.assertEqual(tc._indices(), tg._indices())
self.assertEqual(tc._values(), tg._values())
if t.is_coalesced():
self.assertEqual(tc._indices(), t._indices())
self.assertEqual(tc._values(), t._values())
for empty_i, empty_v, empty_nnz in itertools.product([True, False], repeat=3):
sparse_size = [] if empty_i else [2, 1]
dense_size = [1, 0, 2] if empty_v else [1, 2]
nnz = 0 if empty_nnz else 5
t, _, _ = self._gen_sparse(len(sparse_size), nnz, sparse_size + dense_size, dtype, device, coalesced)
_test_coalesce(t) # this tests correctness
@dtypes(torch.double)
@skipIfTorchDynamo("https://github.com/pytorch/pytorch/issues/89395")
def test_coalesce_reference_cycle(self, device, dtype):
# Test coalesce doesn't create autograd graph cycles (gh-52253)
# Sanity check that the helper class works as expected
t = torch.rand(2)
t_ref = torch._C._WeakTensorRef(t)
self.assertFalse(t_ref.expired())
del t
self.assertTrue(t_ref.expired())
def test_sparse_sum():
i = torch.tensor([[0], [4]], dtype=torch.long, device=device)
v = torch.tensor([[[-0.4567, -1.8797, 0.0380, 1.4316]]],
dtype=dtype, device=device)
S = torch.sparse_coo_tensor(i, v)
S = S.coalesce()
S.requires_grad_(True)
S2 = S.coalesce()
self.assertTrue(S2.is_coalesced())
return torch._C._WeakTensorRef(S2)
ref = test_sparse_sum()
self.assertTrue(ref.expired())
@dtypes(torch.double)
def test_ctor_large_sizes(self, device, dtype):
# Test that integer overflow is detected when computing numel
# of a sparse tensor with large dimensions (gh-57416). Notice
# that numel is computed internally when constructing a
# tensor, hence the overflow may appear during the tensor
# construction step.
N = 100000
indices = torch.tensor([[N, N - 1]] * 4, dtype=torch.int64, device=device)
values = torch.tensor([1, 2], dtype=dtype, device=device)
self.assertRaises(RuntimeError,
lambda: torch.sparse_coo_tensor(
indices, values, (N + 1,) * 4, device=device))
@dtypes(torch.double, torch.cdouble)
def test_ctor_size_checks(self, device, dtype):
indices = self.index_tensor([
[0, 0, 0],
[0, 3, 0],
[0, 0, 0],
[0, 0, 0],
], device=device)
values = torch.tensor([2, 1, 3, 4], dtype=dtype, device=device)
# indices inconsistent with size
self.assertRaises(
RuntimeError,
lambda: self.sparse_tensor(indices, values, torch.Size([2, 1, 1])))
# values inconsistent with size
values = torch.tensor([
[2, 1, 2, 1],
[1, 0, 5, 2],
], dtype=dtype, device=device)
self.assertRaises(
RuntimeError,
lambda: self.sparse_tensor(indices, values, torch.Size([2, 4, 2, 1])))
@coalescedonoff
@dtypes(torch.double)
def test_ctor_is_coalesced_with_gradcheck(self, device, dtype, coalesced):
for sparse_size, nnz in (((3, 3), 5), ((2, 3, 1, 5), 11)):
t, _, _ = self._gen_sparse(len(sparse_size), nnz, sparse_size, dtype, device, coalesced)
self.assertEqual(t.is_coalesced(), coalesced)
def func(indices, values, shape, is_coalesced):
s = torch.sparse_coo_tensor(indices, values, shape, check_invariants=True, is_coalesced=is_coalesced)
self.assertEqual(s.is_coalesced(), is_coalesced)
return s.to_dense(masked_grad=False)
if coalesced:
torch.autograd.gradcheck(func, (t._indices(), t._values().requires_grad_(True), t.shape, False))
torch.autograd.gradcheck(func, (t._indices(), t._values().requires_grad_(True), t.shape, True))
else:
torch.autograd.gradcheck(func, (t._indices(), t._values().requires_grad_(True), t.shape, False))
with self.assertRaisesRegex(RuntimeError,
"cannot set is_coalesced to true if indices correspond to uncoalesced COO tensor"):
torch.autograd.gradcheck(func, (t._indices(), t._values().requires_grad_(True), t.shape, True))
@dtypes(*floating_and_complex_types_and(torch.float16, torch.bfloat16))
@unittest.skipIf(TEST_WITH_CROSSREF, "generator unsupport triggers assertion error")
@gradcheck_semantics()
def test_to_dense_with_gradcheck(self, device, dtype, gradcheck):
def test_tensor(x, res):
x.to_dense() # Tests triple to_dense for memory corruption
x.to_dense()
x.to_dense()
dense_x = x.to_dense()
safe_dense_x = self.safeToDense(x)
dense_x = dense_x.to(res.dtype)
safe_dense_x = safe_dense_x.to(res.dtype)
self.assertEqual(res, dense_x)
self.assertEqual(res, safe_dense_x)
# Only run autograd test for float64
if x.dtype != torch.float64:
return
def fn(x):
return x.to_dense(masked_grad=gradcheck.masked)
x.requires_grad_(True)
gradcheck(fn, (x,))
for value_type in [torch.double, torch.cdouble]:
i = self.index_tensor([
[0, 1, 2, 2],
[0, 0, 0, 3],
[0, 0, 1, 4],
], device=device)
# we don't have to_dense for half types on CPU because it is implemented
# with a slower add_ operation
v = torch.tensor([2, 1, 3, 4], dtype=dtype, device=device)
x = self.sparse_tensor(i, v, torch.Size([3, 4, 5]), dtype=value_type, device=device)
res = torch.tensor([
[[2, 0, 0, 0, 0],
[0, 0, 0, 0, 0],
[0, 0, 0, 0, 0],
[0, 0, 0, 0, 0]],
[[1, 0, 0, 0, 0],
[0, 0, 0, 0, 0],
[0, 0, 0, 0, 0],
[0, 0, 0, 0, 0]],
[[0, 3, 0, 0, 0],
[0, 0, 0, 0, 0],
[0, 0, 0, 0, 0],
[0, 0, 0, 0, 4]],
], dtype=dtype, device=device)
test_tensor(x, res)
test_tensor(res, res)
i = self.index_tensor([
[0, 1, 2, 2],
[0, 0, 0, 3],
[0, 0, 1, 4],
], device=device)
v = torch.empty(4, 0, dtype=dtype, device=device)
x = self.sparse_tensor(i, v, torch.Size([3, 4, 5, 0]), dtype=value_type, device=device)
res = torch.empty((3, 4, 5, 0), dtype=dtype, device=device)
test_tensor(x, res)
@coalescedonoff
@dtypes(torch.float16, torch.bfloat16, torch.float64, torch.int, torch.cfloat, torch.cdouble)
def test_to_sparse(self, device, dtype, coalesced):
shape = [5, 2, 10, 4]
max_nnz = 1
for value_type in [torch.double, torch.cdouble]:
for dim, dim_sz in enumerate(shape, 1):
max_nnz *= dim_sz
rnnz = torch.randint(2, max_nnz, (1,)).item()
for nnz in [0, 1, rnnz]:
expected, _, _ = self._gen_sparse(dim, nnz, shape, dtype=value_type, device=device,
coalesced=coalesced)
expected = expected.to(dtype)
d = expected.to_dense()
result = d.to_sparse(dim)
self.assertEqual(d, result.to_dense())
self.assertEqual(expected.size(), result.size())
self.assertEqual(dim, result.sparse_dim())
@dtypes(torch.double, torch.cdouble)
def test_sparse_bool(self, device, dtype):
a = torch.tensor([True, False], dtype=dtype, device=device).to(torch.bool)
b = a.to_sparse().to_dense()
self.assertEqual(a, b)
@skipIfTorchDynamo("https://github.com/pytorch/pytorch/issues/108667")
@dtypes(torch.double, torch.cdouble)
def test_scalar(self, device, dtype):
# tensor with value
a = self.sparse_tensor(self.index_tensor([], device=device).unsqueeze(1), 12.3, [], dtype=dtype, device=device)
self.assertEqual(1, a._values().numel())
self.assertEqual(a, a.clone())
a_coalesced = a.coalesce()
self.assertTrue(a_coalesced.is_coalesced())
self.assertEqual(torch.tensor(12.3, dtype=dtype, device=device), a.to_dense())
self.assertEqual(a, a.to_dense().to_sparse())
# tensor with multiple values
a = self.sparse_tensor(self.index_tensor([], device=device).unsqueeze(1).expand(0, 2),
[12.3, 12.3], [], dtype=dtype, device=device)
self.assertEqual(2, a._values().numel())
self.assertEqual(a, a.clone())
a_coalesced = a.coalesce()
self.assertTrue(a_coalesced.is_coalesced())
self.assertEqual(torch.tensor(12.3 * 2, dtype=dtype, device=device), a.to_dense())
self.assertEqual(a.coalesce(), a.coalesce().to_dense().to_sparse())
# tensor without value
a = self.sparse_empty((), dtype=dtype, device=device)
self.assertEqual(0, a._values().numel())
self.assertEqual(a, a.clone())
a_coalesced = a.coalesce()
self.assertTrue(a_coalesced.is_coalesced())
self.assertEqual(torch.tensor(0, dtype=dtype, device=device), a.to_dense())
self.assertEqual(a, a.to_dense().to_sparse())
@dtypes(torch.double, torch.cdouble)
def test_shared(self, device, dtype):
i = self.index_tensor([[2]], device=device)
v = torch.tensor([5], dtype=dtype, device=device)
x = self.sparse_tensor(i, v, torch.Size([3]))
v[0] = 6
self.assertEqual(torch.tensor([0, 0, 6], dtype=dtype, device=device), self.safeToDense(x))
i[0][0] = 0
self.assertEqual(torch.tensor([6, 0, 0], dtype=dtype, device=device), self.safeToDense(x))
i = self.index_tensor([[2]], device=device)
v = torch.empty((1, 0), dtype=dtype, device=device)
x = self.sparse_tensor(i, v, torch.Size([3, 0]))
i[0][0] = 0
self.assertEqual(torch.empty((3, 0), dtype=dtype, device=device), self.safeToDense(x))
@dtypes(torch.double, torch.cdouble)
@unittest.skipIf(TEST_WITH_CROSSREF, "generator unsupport triggers assertion error")
@gradcheck_semantics()
def test_to_dense_hybrid(self, device, dtype, gradcheck):
def test_tensor(x, res):
x.to_dense() # Tests double to_dense for memory corruption
x.to_dense()
x.to_dense()
self.assertEqual(res, x.to_dense())
self.assertEqual(res, self.safeToDense(x))
def fn(x):
return x.to_dense(masked_grad=gradcheck.masked)
x.requires_grad_(True)
gradcheck(fn, (x,))
i = self.index_tensor([
[0, 1, 2, 2],
[0, 0, 0, 3],
], device=device)
v = torch.tensor([[2, 3], [1, 2], [3, 4], [4, 5]], dtype=dtype, device=device)
x = self.sparse_tensor(i, v, torch.Size([3, 4, 2]))
res = torch.tensor([
[[2, 3],
[0, 0],
[0, 0],
[0, 0]],
[[1, 2],
[0, 0],
[0, 0],
[0, 0]],
[[3, 4],
[0, 0],
[0, 0],
[4, 5]],
], dtype=dtype, device=device)
test_tensor(x, res)
i = self.index_tensor([
[0, 1, 2, 2],
[0, 0, 0, 3],
], device=device)
v = torch.empty((4, 2, 0), dtype=dtype, device=device)
x = self.sparse_tensor(i, v, torch.Size([3, 4, 2, 0]))
res = torch.empty((3, 4, 2, 0), dtype=dtype, device=device)
test_tensor(x, res)
@dtypes(torch.double, torch.cdouble)
def test_contig(self, device, dtype):
def test_tensor(x, exp_i, exp_v):
x = x.coalesce()
self.assertEqual(exp_i, x._indices())
self.assertEqual(exp_v, x._values())
i = self.index_tensor([
[1, 0, 35, 14, 39, 6, 71, 66, 40, 27],
[92, 31, 62, 50, 22, 65, 89, 74, 56, 34],
], device=device)
v = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8, 9, 10], dtype=dtype, device=device)
x = self.sparse_tensor(i, v, torch.Size([100, 100]))
exp_i = self.index_tensor([
[0, 1, 6, 14, 27, 35, 39, 40, 66, 71],
[31, 92, 65, 50, 34, 62, 22, 56, 74, 89],
], device=device)
exp_v = torch.tensor([2, 1, 6, 4, 10, 3, 5, 9, 8, 7], dtype=dtype, device=device)
test_tensor(x, exp_i, exp_v)
i = self.index_tensor([
[2, 0, 2, 1],
[0, 0, 3, 0],
[1, 0, 4, 0],
], device=device)
v = torch.tensor([3, 2, 4, 1], dtype=dtype, device=device)
x = self.sparse_tensor(i, v, torch.Size([3, 4, 5]))
exp_i = self.index_tensor([
[0, 1, 2, 2],
[0, 0, 0, 3],
[0, 0, 1, 4],
], device=device)
exp_v = torch.tensor([2, 1, 3, 4], dtype=dtype, device=device)
test_tensor(x, exp_i, exp_v)
i = self.index_tensor([
[2, 0, 2, 1],
[0, 0, 3, 0],
[1, 0, 4, 0],
], device=device)
v = torch.empty([4, 0], dtype=dtype, device=device)
x = self.sparse_tensor(i, v, torch.Size([3, 4, 5, 0]))
exp_i = self.index_tensor([
[0, 1, 2, 2],
[0, 0, 0, 3],
[0, 0, 1, 4],
], device=device)
exp_v = torch.empty([4, 0], dtype=dtype, device=device)
test_tensor(x, exp_i, exp_v)
# Duplicate indices
i = self.index_tensor([
[0, 0, 2, 0],
[0, 0, 3, 0],
[0, 0, 4, 0],
], device=device)
v = torch.tensor([3, 2, 4, 1], dtype=dtype, device=device)
x = self.sparse_tensor(i, v, torch.Size([3, 4, 5]))
exp_i = self.index_tensor([
[0, 2],
[0, 3],
[0, 4],
], device=device)
exp_v = torch.tensor([6, 4], dtype=dtype, device=device)
test_tensor(x, exp_i, exp_v)
i = self.index_tensor([
[0, 0, 2, 0],
[0, 0, 3, 0],
[0, 0, 4, 0],
], device=device)
v = torch.empty([4, 0], dtype=dtype, device=device)
x = self.sparse_tensor(i, v, torch.Size([3, 4, 5, 0]))
exp_i = self.index_tensor([
[0, 2],
[0, 3],
[0, 4],
], device=device)
exp_v = torch.empty([2, 0], dtype=dtype, device=device)
test_tensor(x, exp_i, exp_v)
@dtypes(torch.double, torch.cdouble)
def test_contig_hybrid(self, device, dtype):
def test_tensor(x, exp_i, exp_v):
x = x.coalesce()
self.assertEqual(exp_i, x._indices())
self.assertEqual(exp_v, x._values())
i = self.index_tensor([
[1, 0, 35, 14, 39, 6, 71, 66, 40, 27],
[92, 31, 62, 50, 22, 65, 89, 74, 56, 34],
], device=device)
v = torch.tensor([
[1, 2], [2, 3], [3, 4], [4, 5], [5, 6],
[6, 7], [7, 8], [8, 9], [9, 10], [10, 11],
], dtype=dtype, device=device)
x = self.sparse_tensor(i, v, torch.Size([100, 100, 2]))
exp_i = self.index_tensor([
[0, 1, 6, 14, 27, 35, 39, 40, 66, 71],
[31, 92, 65, 50, 34, 62, 22, 56, 74, 89],
], device=device)
exp_v = torch.tensor([
[2, 3], [1, 2], [6, 7], [4, 5], [10, 11],
[3, 4], [5, 6], [9, 10], [8, 9], [7, 8],
], dtype=dtype, device=device)
test_tensor(x, exp_i, exp_v)
i = self.index_tensor([
[2, 0, 2, 1],
[0, 0, 3, 0],
[1, 0, 4, 0],
], device=device)
v = torch.tensor([[3, 3, 3], [2, 2, 2], [4, 4, 4], [1, 1, 1]], dtype=dtype, device=device)
x = self.sparse_tensor(i, v, torch.Size([3, 4, 5, 3]))
exp_i = self.index_tensor([
[0, 1, 2, 2],
[0, 0, 0, 3],
[0, 0, 1, 4],
], device=device)
exp_v = torch.tensor([[2, 2, 2], [1, 1, 1], [3, 3, 3], [4, 4, 4]], dtype=dtype, device=device)
test_tensor(x, exp_i, exp_v)
i = self.index_tensor([
[2, 0, 2, 1],
[0, 0, 3, 0],
[1, 0, 4, 0],
], device=device)
v = torch.empty([4, 3, 0], dtype=dtype, device=device)
x = self.sparse_tensor(i, v, torch.Size([3, 4, 5, 3, 0]))
exp_i = self.index_tensor([
[0, 1, 2, 2],
[0, 0, 0, 3],
[0, 0, 1, 4],
], device=device)
exp_v = torch.empty([4, 3, 0], dtype=dtype, device=device)
test_tensor(x, exp_i, exp_v)
# Duplicate indices
i = self.index_tensor([
[0, 0, 2, 0],
[0, 0, 3, 0],
[0, 0, 4, 0],
], device=device)
v = torch.tensor([[3, 2, 3], [2, 1, 1], [4, 3, 4], [1, 1, 1]], dtype=dtype, device=device)
x = self.sparse_tensor(i, v, torch.Size([3, 4, 5, 3]))
exp_i = self.index_tensor([
[0, 2],
[0, 3],
[0, 4],
], device=device)
exp_v = torch.tensor([[6, 4, 5], [4, 3, 4]], dtype=dtype, device=device)
test_tensor(x, exp_i, exp_v)
i = self.index_tensor([
[0, 0, 2, 0],
[0, 0, 3, 0],
[0, 0, 4, 0],
], device=device)
v = torch.empty([4, 3, 0], dtype=dtype, device=device)
x = self.sparse_tensor(i, v, torch.Size([3, 4, 5, 3, 0]))
exp_i = self.index_tensor([
[0, 2],
[0, 3],
[0, 4],
], device=device)
exp_v = torch.empty([2, 3, 0], dtype=dtype, device=device)
test_tensor(x, exp_i, exp_v)
@coalescedonoff
@dtypes(torch.double, torch.cdouble)
def test_clone(self, device, dtype, coalesced):
def test_shape(sparse_dims, nnz, with_size):
x = self._gen_sparse(sparse_dims, nnz, with_size, dtype, device, coalesced)[0]
if not coalesced:
self.assertFalse(x.is_coalesced())
y = x.clone()
self.assertFalse(y.is_coalesced())
x = x.coalesce()
self.assertTrue(x.is_coalesced())
y = x.clone()
self.assertTrue(y.is_coalesced())
test_shape(4, 20, 5)
test_shape(3, 10, [100, 100, 100, 5, 5, 5, 0])
test_shape(3, 0, [0, 0, 100, 5, 5, 5, 0])
@coalescedonoff
@dtypes(torch.double, torch.cdouble, torch.bfloat16)
@precisionOverride({torch.bfloat16: 2e-2})
def test_Sparse_to_Sparse_copy_(self, device, dtype, coalesced):
# This is for testing torch.copy_(SparseTensor, SparseTensor)
sparse_dims = 3
nnz = 10
sizes = [2, 3, 4, 5] # hybrid sparse
x1, _, _ = self._gen_sparse(sparse_dims, nnz, sizes, dtype, device, coalesced)
x2, _, _ = self._gen_sparse(sparse_dims, nnz + 10, sizes, dtype, device, coalesced)
# test copy
x2_dense = x2.to_dense()
x1.copy_(x2)
self.assertEqual(x2_dense, x1.to_dense())
# test type conversion (when x1.copy_(x2), x1.dtype should stay the same)
x1 = x1.to(torch.float32)
x2 = x2.to(torch.float16)
x1_dtype = x1.dtype
x1.copy_(x2)
self.assertEqual(x1_dtype, x1.dtype)
x2 = x2.to(torch.float64)
x1_dtype = x1.dtype
x1.copy_(x2)
self.assertEqual(x1_dtype, x1.dtype)
# test no broadcast
self.assertRaises(RuntimeError, lambda: x1.copy_(x2.narrow_copy(0, 0, 1)))
# test raise error on copy_() between dense and sparse Tensors
self.assertRaises(RuntimeError, lambda: x1.copy_(torch.randn(5, 5)))
# test autograd
x1, _, _ = self._gen_sparse(sparse_dims, nnz, sizes, dtype, device, coalesced)
x2, _, _ = self._gen_sparse(sparse_dims, nnz + 10, sizes, dtype, device, coalesced)
x2.requires_grad_(True)
x1.copy_(x2)
y = x1 * 2
x2_clone = x2.clone()
y.backward(x2_clone)
expected_grad = x2_clone * 2
self.assertEqual(expected_grad.to_dense(), x2.grad.to_dense())
self.assertEqual(None, x1.grad)
@coalescedonoff
@unittest.skipIf(not TEST_MULTIGPU, "multi-GPU not supported")
@dtypes(torch.double, torch.cdouble)
def test_Sparse_to_Sparse_copy_multi_gpu(self, device, dtype, coalesced):
# This is for testing torch.copy_(SparseTensor, SparseTensor) across GPU devices
sparse_dims = 3
nnz = 10
sizes = [2, 3, 4, 5] # hybrid sparse
x1, _, _ = self._gen_sparse(sparse_dims, nnz, sizes, dtype, device, coalesced)
x2, _, _ = self._gen_sparse(sparse_dims, nnz + 10, sizes, dtype, device, coalesced)
x1 = x1.to('cuda:0')
def test_cross_device(x1, x2):
x1_device = x1.device
x1.copy_(x2)
self.assertEqual(x2.to('cuda:0').to_dense(), x1.to_dense())
self.assertEqual(x1_device, x1.device)
test_cross_device(x1, x2.to('cuda:1')) # test across gpu devices
test_cross_device(x1, x2.to('cpu')) # test between cpu and gpu
# test autograd
x2 = x2.to('cuda:1')
x2.requires_grad_(True)
x1.copy_(x2)
y = x1 * 2
x2_clone = x2.clone().to('cuda:0')
y.backward(x2_clone)
expected_grad = x2_clone * 2
self.assertEqual(expected_grad.to_dense(), x2.grad.to('cuda:0').to_dense())
self.assertEqual(None, x1.grad)
@onlyCUDA
def test_cuda_empty(self, device):
def test_tensor(x):
y = x.to(device)
self.assertEqual(x.sparse_dim(), y.sparse_dim())
self.assertEqual(x.dense_dim(), y.dense_dim())
x = y.cpu()
self.assertEqual(y.sparse_dim(), x.sparse_dim())
self.assertEqual(y.dense_dim(), x.dense_dim())
x = torch.sparse_coo_tensor((2, 3, 4), dtype=torch.float32)
test_tensor(x)
x = torch.sparse_coo_tensor((2, 3, 4), dtype=torch.float16)
test_tensor(x)
x = torch.sparse_coo_tensor((2, 3, 4), dtype=torch.float16)
test_tensor(x)
x = torch.sparse_coo_tensor((2, 3, 4, 0), dtype=torch.float32)
test_tensor(x)
@coalescedonoff
@dtypes(torch.double, torch.cdouble)
def test_transpose(self, device, dtype, coalesced):
def test_shape(sparse_dims, nnz, with_size):
x = self._gen_sparse(sparse_dims, nnz, with_size, dtype, device, coalesced)[0]
y = self.safeToDense(x)
for i, j in itertools.combinations(range(4), 2):
x = x.transpose_(i, j)
y = y.transpose(i, j)
self.assertEqual(self.safeToDense(x), y)
x = x.transpose(i, j)
y = y.transpose(i, j)
self.assertEqual(self.safeToDense(x), y)
test_shape(4, 6, 3)
test_shape(4, 3, [7, 7, 7, 3, 3, 3, 0])
test_shape(4, 0, [0, 0, 7, 3, 3, 3, 0])
@coalescedonoff
@dtypes(torch.double, torch.cdouble)
@unittest.skipIf(TEST_WITH_CROSSREF, "generator unsupport triggers assertion error")
@gradcheck_semantics()
def test_permute(self, device, dtype, coalesced, gradcheck):
# trivial checks
s = torch.rand(3, 3, 3, device=device, dtype=dtype).to_sparse()
with self.assertRaisesRegex(RuntimeError, "does not match the length"):
s.permute(dims=(1, 0))
with self.assertRaisesRegex(RuntimeError, "duplicate dims"):
s.permute(dims=(1, 1, 1))
# Calling permute on a sparse tensor with an empty tuple used to segfault,
# see https://github.com/pytorch/pytorch/issues/116325
x = torch.rand((), device=device, dtype=dtype).to_sparse()
x.permute(())
self.assertEqual(len(x.values()), 1)
def test_shape(sparse_dims, nnz, with_size):
ndim = len(with_size)
valid_sparse_dims = torch.arange(-ndim, -ndim + sparse_dims)
valid_dense_dims = torch.arange(-ndim + sparse_dims, 0)
for dims in itertools.permutations(range(-ndim, 0)):
s = self._gen_sparse(sparse_dims, nnz, with_size, dtype, device, coalesced)[0]
d = self.safeToDense(s)
dims_sparse, _ = torch.tensor(dims[:sparse_dims]).sort()
dims_dense, _ = torch.tensor(dims[sparse_dims:]).sort()
if (valid_sparse_dims == dims_sparse).all() and (valid_dense_dims == dims_dense).all():
# if valid permutation, test for correctness
s_permuted = s.permute(dims)
self.assertEqual(s_permuted, d.permute(dims))
# if s is coalesced, and perm does not touch 0-dim,
# the result has to be coalesced as well
if dims[0] == 0:
self.assertEqual(s_permuted.is_coalesced(), s.is_coalesced())
else:
self.assertFalse(s_permuted.is_coalesced())
gradcheck(lambda t: t.permute(dims).to_dense(masked_grad=gradcheck.masked), s.requires_grad_())
else:
# otherwise check if exception is thrown
fail_message = "transpositions between sparse and dense dimensions are not allowed"
with self.assertRaisesRegex(RuntimeError, fail_message):
s.permute(dims)
test_shape(2, 3, [2, 3, 4, 5])
test_shape(2, 3, [2, 2, 0])
# if nnz=0, it is not true that t == t.to_dense().to_sparse()
# unless t.sparse_dim == t.dim (i.e. t is not hybrid)
test_shape(3, 0, [0, 0, 2])
@coalescedonoff
@onlyCPU
@dtypes(torch.double)
def test_coalesce_transpose_mm(self, device, dtype, coalesced):
def test_shape(di, dj, dk, nnz):
x, _, _ = self._gen_sparse(2, nnz, [dj, di], dtype, device, coalesced)
y = torch.randn(dj, dk, dtype=dtype, device=device)
x_coalesced = x.coalesce()
self.assertTrue(x_coalesced.is_coalesced())
x_coalesced_t = x_coalesced.t()
# Transpose is `colasced`-preserving if the indices tensor is empty.
self.assertEqual(x_coalesced_t.is_coalesced(), di * nnz == 0)
res = torch.mm(x_coalesced_t, y)
expected = torch.mm(self.safeToDense(x_coalesced_t), y)
self.assertEqual(res, expected)
test_shape(10, 20, 30, 20)
test_shape(0, 20, 30, 0)
test_shape(10, 0, 30, 0)
test_shape(10, 20, 0, 0)
test_shape(10, 20, 0, 20)
@skipIfTorchDynamo("https://github.com/pytorch/torchdynamo/issues/1166")
@dtypes(torch.double, torch.cdouble)
def test_t_empty(self, device, dtype):
def test_in_place(x):
shape_original = x.shape
x.t_()
self.assertEqual(torch.Size([shape_original[1], shape_original[0]]), x.size())
self.assertEqual(0, x._indices().numel())
self.assertEqual(0, x._values().numel())
self.assertEqual(x.sparse_dim(), 2)
self.assertEqual(x.dense_dim(), 0)
def test_not_in_place(x):
shape_original = x.shape
y = x.t()
self.assertEqual(torch.Size([shape_original[1], shape_original[0]]), y.size())
self.assertEqual(0, y._indices().numel())
self.assertEqual(0, y._values().numel())
self.assertEqual(x.sparse_dim(), 2)
self.assertEqual(x.dense_dim(), 0)
x = self.sparse_empty(2, 3, dtype=dtype, device=device)
test_in_place(x)
test_not_in_place(x)
x = self.sparse_empty(2, 0, dtype=dtype, device=device)
test_in_place(x)
test_not_in_place(x)
@coalescedonoff
@dtypes(torch.double, torch.cdouble)
def test_add_zeros(self, device, dtype, coalesced):
def test_shape(sparse_dims, nnz, sizes):
x, _, _ = self._gen_sparse(sparse_dims, nnz, sizes, dtype, device, coalesced)
zeros = torch.sparse_coo_tensor(sizes, device=x.device)
r1 = zeros + x
r2 = x + zeros
self.assertEqual(r1, x)
self.assertEqual(r2, x)
test_shape(1, 20, [1])
test_shape(4, 20, [3, 17, 19, 5])
test_shape(2, 20, [3, 17, 19, 5])
test_shape(2, 20, [3, 17, 19, 0])
@dtypes(torch.double, torch.cdouble)
def test_add_sub_nnz(self, device, dtype):
# nnz should not grow unbounded (gh-34964)
x = torch.randn(10, dtype=dtype, device=device).to_sparse()
x.add_(x)
x.add_(x)
self.assertLessEqual(x._nnz(), 10)
x.sub_(2 * x)
x.sub_(2 * x)
self.assertLessEqual(x._nnz(), 10)
@coalescedonoff
@dtypes(torch.double, torch.cdouble)
def test_cat(self, device, dtype, coalesced):
# shapes: list of tuples (sparse_dims, nnz, sizes)
def test_shapes(shapes, dim, fail_message=None):
inputs = [self._gen_sparse(shape[0], shape[1], shape[2], dtype, device, coalesced)[0]
for shape in shapes]
if fail_message:
with self.assertRaisesRegex(RuntimeError, fail_message):
torch.cat(inputs, dim)
else:
result = torch.cat(inputs, dim)
dense_result = torch.cat([t.to_dense() for t in inputs], dim)
self.assertEqual(dense_result, result.to_dense())
test_shapes(
[(3, 10, [2, 3, 4]), (3, 10, [2, 1, 4]), (3, 10, [2, 4, 4])], 1)
# mismatched sizes
test_shapes([(3, 10, [2, 3, 4]), (3, 10, [2, 1, 4])], 0,
"All tensors must have the same shape: \\[2, 3, 4].*\\[2, 1, 4]")
# hybrid sparse/dense
test_shapes(
[(2, 10, [2, 3, 4]), (2, 10, [2, 1, 4]), (2, 10, [2, 4, 4])], 1)
# cat along dense dim
test_shapes([(2, 10, [2, 3, 4]), (2, 10, [2, 3, 7])], 2)
test_shapes([(1, 10, [2, 3, 4]), (1, 10, [2, 3, 4])], 1)
test_shapes([(1, 10, [2, 3, 4]), (1, 10, [2, 3, 4])], 2)
# mismatched dimensions
test_shapes([(2, 10, [2, 3, 4]), (3, 10, [2, 3, 4])], 0,
"All tensors must have the same.*2, 1, but tensor at position 1 has 3, 0.")
# wrapped dimension
test_shapes(
[(3, 10, [2, 3, 4]), (3, 10, [2, 1, 4]), (3, 10, [2, 4, 4])], -2)
# sparse with dense
sp = self._gen_sparse(3, 10, [2, 3, 4], dtype, device, coalesced)[0]
dn = sp.to_dense()
with self.assertRaisesRegex(RuntimeError,
"Concatenating sparse tensors, but a dense tensor was found at position 1."):
torch.cat((sp, dn))
@coalescedonoff
@dtypes(torch.double, torch.cdouble)
def test_unsqueeze(self, device, dtype, coalesced):
def test_shape(sparse_dims, nnz, sizes, unsqueeze_dim, fail_message=None):
x, _, _ = self._gen_sparse(sparse_dims, nnz, sizes, dtype, device, coalesced)
if fail_message:
with self.assertRaisesRegex(IndexError, fail_message):
torch.unsqueeze(x, unsqueeze_dim)
else:
result = torch.unsqueeze(x, unsqueeze_dim)
dense_result = torch.unsqueeze(x.to_dense(), unsqueeze_dim)
self.assertEqual(dense_result, result.to_dense())
# basic case
test_shape(3, 10, [5, 7, 11], 0)
# hybrid sparse/dense, unsqueeze along sparse dim
test_shape(3, 10, [5, 7, 11, 13, 17], 0)
test_shape(3, 10, [5, 7, 11, 13, 17], 3)
# unsqueeze along dense dimensions
test_shape(3, 10, [5, 7, 11, 13, 17], 4)
test_shape(3, 10, [5, 7, 11, 13, 17], 5)
# wrapped dimensions
test_shape(3, 10, [5, 7, 11, 13, 17], -1)
test_shape(3, 10, [5, 7, 11, 13, 17], -6)
# bounds
test_shape(3, 10, [5, 7, 11, 13, 17], -7, "Dimension out of range")
test_shape(3, 10, [5, 7, 11, 13, 17], 6, "Dimension out of range")
@coalescedonoff
@dtypes(torch.double, torch.cdouble)
def test_select(self, device, dtype, coalesced):
def test_shape(sparse_dims, nnz, sizes, select_dim, select_index, fail_message=None):
x, _, _ = self._gen_sparse(sparse_dims, nnz, sizes, dtype, device, coalesced)
if fail_message:
with self.assertRaisesRegex(IndexError, fail_message):
torch.select(x, select_dim, select_index)
else:
result = torch.select(x, select_dim, select_index)
if result.is_sparse:
result = result.to_dense()
dense_result = torch.select(x.to_dense(), select_dim, select_index)
self.assertEqual(dense_result, result)
sizes = [5, 7, 11, 13, 17]
# hybrid sparse/dense, select sparse dim, result is dense
for i in range(sizes[0]):
test_shape(1, 10, sizes, 0, i)
test_shape(1, 10, sizes, 0, sizes[0] + 1, r'select[(][)][:] index \d out of range.*')
# hybrid sparse/dense, select sparse dim, result is sparse
for d in range(3):
for i in range(sizes[d]):
test_shape(3, 10, sizes, d, i)
# hybrid sparse/dense, select dense dim, result is sparse
for d in range(1, 3):
for i in range(sizes[d]):
test_shape(1, 10, sizes, d, i)
@dtypes(*integral_types())
def test_select_no_type_promotion(self, device, dtype):
# see https://github.com/pytorch/pytorch/issues/82150
idx = torch.tensor([[0, 0, 0, 1, 1, 1], [0, 0, 0, 1, 1, 1]])
val = torch.ones(6, dtype=dtype)
s = torch.sparse_coo_tensor(idx, val, size=(3, 3))
for t in (s, s * torch.tensor(0, dtype=dtype)):
# empty checks
self.assertEqual(t.dtype, t[2].dtype)
self.assertEqual(t.dtype, t[0, 1].dtype)
# sum should not promote
self.assertEqual(t.dtype, t[0, 0].dtype)
self.assertEqual(t.dtype, t[1, 1].dtype)
@coalescedonoff
@dtypes(torch.double, torch.cdouble)
def test_index_select(self, device, dtype, coalesced):
def test_shape(sparse_dims, nnz, sizes, select_dim, select_index, fail_message=None):
if isinstance(select_index, int):
select_index = [select_index]
if isinstance(select_index, list):
select_index = torch.tensor(select_index, device=device, dtype=torch.long)
x, _, _ = self._gen_sparse(sparse_dims, nnz, sizes, dtype, device, coalesced)
if fail_message:
with self.assertRaisesRegex(IndexError, fail_message):
torch.index_select(x, select_dim, select_index)
else:
result = torch.index_select(x, select_dim, select_index)
if result.is_sparse:
result = result.to_dense()
dense_result = torch.index_select(x.to_dense(), select_dim, select_index)
self.assertEqual(dense_result, result)
sizes = [5, 7, 11, 13, 17]
for d in range(len(sizes)):
for index in [0, sizes[d] - 1, [0, sizes[d] // 2, sizes[d] - 1]]:
test_shape(1, 10, sizes, d, index)
test_shape(len(sizes) // 2, 10, sizes, d, index)
test_shape(len(sizes), 10, sizes, d, index)
def _test_index_select_exhaustive_index(self, sizes, dims, device, dtype, coalesced):
t = make_tensor(sizes, dtype=dtype, device=device)
t_sparse = t.to_sparse().coalesce() if coalesced else t.to_sparse()
t_small_sparse, _, _ = self._gen_sparse(len(sizes), 2, sizes, dtype, device, coalesced)
t_small = t_small_sparse.to_dense()
for d in dims:
# NOTE: indices are negative
idx_dim_d_range = list(range(-sizes[d], 0))
for idx_len in range(sizes[d], sizes[d] + 1):
# creates all possible valid indices into dim d of lenght idx_len
for idx in itertools.product(*itertools.repeat(idx_dim_d_range, idx_len)):
t_idx = torch.tensor(idx, dtype=torch.long, device=device)
# NOTE: index_select for dense does not support negative indices,
# hence + sizes[d]. See https://github.com/pytorch/pytorch/issues/76347
# tests the nnz > sizes[d] branch
dense_result = t.index_select(d, t_idx + sizes[d])
sparse_result = t_sparse.index_select(d, t_idx)
self.assertEqual(dense_result, sparse_result)
# tests the nnz <= sizes[d] branch
small_dense_result = t_small.index_select(d, t_idx + sizes[d])
small_sparse_result = t_small_sparse.index_select(d, t_idx)
self.assertEqual(small_dense_result, small_sparse_result)
@coalescedonoff
@dtypes(torch.double, torch.cdouble)
def test_index_select_exhaustive_index_small(self, device, dtype, coalesced):
# will trigger brute-force algo
self._test_index_select_exhaustive_index((3, 3, 4), range(3), device, dtype, coalesced)
@coalescedonoff
@dtypes(torch.double, torch.cdouble)
def test_index_select_exhaustive_index_large(self, device, dtype, coalesced):
# will trigger more sophisticated algos
self._test_index_select_exhaustive_index((100, 50, 3, 3), (2, 3), device, dtype, coalesced)
@coalescedonoff
@dtypes(torch.double, torch.cdouble)
def test_index_select_empty_and_non_contiguous_index(self, device, dtype, coalesced):
# empty index
idx_empty = torch.tensor([], dtype=torch.long, device=device)
t = make_tensor((5, 5), dtype=dtype, device=device)
res_dense = t.index_select(0, idx_empty)
res_sparse = t.to_sparse().index_select(0, idx_empty)
self.assertEqual(res_dense, res_sparse)
# non-contigous index
idx = torch.randint(low=0, high=5, size=(10, 2), device=device)[:, 0]
def run_test(sizes):
# case nnz > size[d]
t = make_tensor(sizes, dtype=dtype, device=device)
res_dense = t.index_select(0, idx)
res_sparse = t.to_sparse().index_select(0, idx)
self.assertEqual(res_dense, res_sparse)
# case nnz <= size[d]
t_small_sparse, _, _ = self._gen_sparse(len(sizes), 2, sizes, dtype, device, coalesced)
res_sparse = t_small_sparse.index_select(0, idx)
res_dense = t_small_sparse.to_dense().index_select(0, idx)
self.assertEqual(res_dense, res_sparse)
# brute-force
run_test((10, 10))
# more sophisticated algos
run_test((10, 100, 100))
@onlyCPU
@coalescedonoff
@dtypes(torch.double, torch.cdouble)
def test_index_select_parallelization(self, device, dtype, coalesced):
"""
Test with sizes that will trigger parallelization (i.e. with sizes
that are >= at::internal::GRAIN_SIZE)
"""
def run_test(nnz, size):
t_sparse, _, _ = self._gen_sparse(1, nnz, (size,), dtype, device, coalesced)
t_dense = t_sparse.to_dense()
# idx_small to (sort) and (binary) search into t_sparse
idx_small = torch.randint(size, (nnz // 2,), device=device)
# idx_large to (sort) and (binary) search into idx_large
# NOTE: when coalesced=True, the (binary) search will be
# done over t_sparse anyway, as it is already sorted.
idx_large = torch.randint(size, (nnz * 2,), device=device)
for idx in (idx_small, idx_large):
res_dense = t_dense.index_select(0, idx)
res_sparse = t_sparse.index_select(0, idx)
self.assertEqual(res_dense, res_sparse)
# NOTE: GRAIN_SIZE = 32768
# case nnz <= size[d]
tlen = 70000 # > 2 * GRAIN_SIZE
run_test(tlen, tlen)
# case nnz > size[d]
run_test(tlen, tlen // 2)
@onlyCPU
@coalescedonoff
@dtypes(torch.double, torch.cdouble)
def test_mm(self, device, dtype, coalesced):
def test_shape(di, dj, dk, nnz):
x, _, _ = self._gen_sparse(2, nnz, [di, dj], dtype, device, coalesced)
t = torch.randn(di, dk, dtype=dtype, device=device)
y = torch.randn(dj, dk, dtype=dtype, device=device)
alpha = random.random()
beta = random.random()
res = torch.addmm(t, x, y, beta=beta, alpha=alpha)
expected = torch.addmm(t, self.safeToDense(x), y, beta=beta, alpha=alpha)
self.assertEqual(res, expected)
res = torch.addmm(t, x, y)
expected = torch.addmm(t, self.safeToDense(x), y)
self.assertEqual(res, expected)
res = torch.mm(x, y)
expected = torch.mm(self.safeToDense(x), y)
self.assertEqual(res, expected)
test_shape(10, 100, 100, 20)
test_shape(100, 1000, 200, 20)
test_shape(64, 10000, 300, 20)
test_shape(0, 100, 100, 0)
test_shape(10, 0, 100, 0)
test_shape(10, 100, 0, 0)
test_shape(10, 100, 0, 20)
@unittest.skipIf(
IS_WINDOWS and TEST_CUDA,
"bmm sparse-dense CUDA is not yet supported in Windows, at least up to CUDA 10.1"
)
@coalescedonoff
@dtypes(torch.double)
def test_bmm(self, device, dtype, coalesced):
def test_shape(num_mats, dim_i, dim_j, dim_k, nnz):
a_list = []
b_list = []
for mat_idx in range(num_mats):
a_mat = self._gen_sparse(2, nnz, [dim_i, dim_j], dtype, device, coalesced)[0]
b_mat = torch.randn([dim_j, dim_k], dtype=dtype, device=device)
a_list.append(a_mat)
b_list.append(b_mat)
a = torch.stack(a_list)
b = torch.stack(b_list)
ab = a.bmm(b)
# Compare each matrix against result from mm()
for mat_idx in range(num_mats):
a_mat = a_list[mat_idx]
b_mat = b_list[mat_idx]
ab_mat_bmm = ab[mat_idx]
ab_mat_mm = a_mat.mm(b_mat)
self.assertEqual(ab_mat_bmm, ab_mat_mm)
test_shape(10, 10, 100, 99, 20)
test_shape(10, 100, 1000, 200, 20)
test_shape(10, 64, 10000, 300, 20)
test_shape(10, 0, 100, 99, 0)
test_shape(10, 10, 0, 100, 0)
test_shape(10, 10, 100, 0, 0)
test_shape(10, 10, 100, 0, 20)
test_shape(10, 10, 100, 0, 20)
a = torch.rand([10, 23, 32], dtype=dtype, device=device)
a[3] = torch.zeros(23, 32, dtype=dtype, device=device)
a[6] = torch.zeros(23, 32, dtype=dtype, device=device)
a = a.to_sparse()
b = torch.rand([10, 32, 10], dtype=dtype, device=device)
b[4] = torch.zeros(32, 10, dtype=dtype, device=device)
b[6] = torch.zeros(32, 10, dtype=dtype, device=device)
ab = a.bmm(b)
for mat_idx in range(ab.size(0)):
ab_mat = ab[mat_idx]
ab_mat_check = a[mat_idx].mm(b[mat_idx])
self.assertEqual(ab_mat, ab_mat_check)
ab_traspose_check = b.transpose(1, 2).to_sparse().bmm(
a.transpose(1, 2).to_dense()
).transpose(1, 2)
self.assertEqual(ab, ab_traspose_check)
@onlyCUDA
@coalescedonoff
@dtypes(torch.double)
@unittest.skipIf(
IS_WINDOWS,
"bmm sparse-dense CUDA is not yet supported in Windows, at least up to CUDA 10.1"
)
def test_bmm_deterministic(self, device, dtype, coalesced):
def test_shape(num_mats, dim_i, dim_j, dim_k, nnz):
a_list = []
b_list = []
for mat_idx in range(num_mats):
a_list.append(self._gen_sparse(2, nnz, [dim_i, dim_j], dtype, device, coalesced)[0])
b_list.append(torch.randn([dim_j, dim_k], dtype=dtype, device=device))
a = torch.stack(a_list).cuda()
b = torch.stack(b_list).cuda()
with DeterministicGuard(torch.are_deterministic_algorithms_enabled()):
torch.use_deterministic_algorithms(False)
ab_nondeterministic = torch.bmm(a, b)
torch.use_deterministic_algorithms(True)
ab_deterministic = torch.bmm(a, b)
diff_abs = (ab_deterministic - ab_nondeterministic).abs()
diff_rel = diff_abs / ab_deterministic.abs()
diff_rel[torch.isnan(diff_rel)] = 0
# deterministic and non-deterministic results should either be
# equal or within a small relative difference
equal_abs_or_rel = diff_abs.eq(0).logical_or(diff_rel.lt(0.001))
self.assertTrue(equal_abs_or_rel.all())
test_shape(10, 10, 100, 99, 20)
test_shape(10, 100, 1000, 200, 20)
test_shape(10, 64, 10000, 300, 20)
test_shape(10, 0, 100, 99, 0)
test_shape(10, 10, 0, 100, 0)
test_shape(10, 10, 100, 0, 0)
test_shape(10, 10, 100, 0, 20)
test_shape(10, 10, 100, 0, 20)
@onlyCUDA
@unittest.skipIf(
not IS_WINDOWS or not TEST_WITH_ROCM,
"this test ensures bmm sparse-dense CUDA gives an error when run on Windows with CUDA < 11.0"
)
@dtypes(torch.double)
def test_bmm_windows_error(self, device, dtype):
a = torch.rand(2, 2, 2, dtype=dtype).to_sparse().cuda()
b = torch.rand(2, 2, 2, dtype=dtype).cuda()
with self.assertRaisesRegex(
RuntimeError,
"bmm sparse-dense CUDA is not supported on Windows with cuda before 11.0"):
ab = a.bmm(b)
@onlyCPU
@coalescedonoff
@dtypes(torch.double, torch.cdouble)
def test_saddmm(self, device, dtype, coalesced):
def test_shape(di, dj, dk, nnz):
x = self._gen_sparse(2, nnz, [di, dj], dtype, device, coalesced)[0]
t = self._gen_sparse(2, nnz, [di, dk], dtype, device, coalesced)[0]
y = torch.randn(dj, dk, dtype=dtype, device=device)
alpha = random.random()
beta = random.random()
res = torch.saddmm(t, x, y, beta=beta, alpha=alpha)
expected = torch.addmm(self.safeToDense(t), self.safeToDense(x), y, beta=beta, alpha=alpha)
self.assertEqual(self.safeToDense(res), expected)
res = torch.saddmm(t, x, y)
expected = torch.addmm(self.safeToDense(t), self.safeToDense(x), y)
self.assertEqual(self.safeToDense(res), expected)
res = torch.smm(x, y)
expected = torch.mm(self.safeToDense(x), y)
self.assertEqual(self.safeToDense(res), expected)
test_shape(7, 5, 3, 20)
test_shape(1000, 100, 100, 20)
test_shape(3000, 64, 300, 20)
test_shape(0, 100, 100, 0)
test_shape(1000, 0, 100, 0)
test_shape(1000, 100, 0, 0)
@onlyCPU
@coalescedonoff
# adding a graph break before self.assertFalse(weight._indices().is_contiguous())
# makes the test pass so some existent sparse related bug
@skipIfTorchDynamo("skip")
@dtypes(torch.double, torch.cdouble)
def test_sspaddmm(self, device, dtype, coalesced):
def test_shape(di, dj, dk, nnz):
x = self._gen_sparse(2, nnz, [di, dj], dtype, device, coalesced)[0]
t = self._gen_sparse(2, nnz, [di, dk], dtype, device, coalesced)[0]
y = torch.randn(dj, dk, dtype=dtype, device=device)
alpha = random.random()
beta = random.random()
res = t.sspaddmm(x, y, beta=beta, alpha=alpha)
expected = torch.addmm(self.safeToDense(t), self.safeToDense(x), y, beta=beta, alpha=alpha)
self.assertEqual(self.safeToDense(res), expected)
res = t.sspaddmm(x, y)
expected = torch.addmm(self.safeToDense(t), self.safeToDense(x), y)
self.assertEqual(self.safeToDense(res), expected)
test_shape(7, 5, 3, 20)
test_shape(1000, 100, 100, 20)
test_shape(3000, 64, 300, 20)
test_shape(0, 100, 100, 0)
test_shape(1000, 0, 100, 0)
test_shape(1000, 100, 0, 0)
# Test code from issue https://github.com/pytorch/pytorch/issues/45113
batch_size, input_size, hidden_size = 5, 3, 7
# Create coalesced sparse tensor with non-contiguous indices
weight = torch.randn(hidden_size, input_size, dtype=dtype, device=device).to_sparse()
self.assertTrue(weight.is_coalesced())
non_contig_indices = weight.indices().mT.contiguous().mT
weight = torch.sparse_coo_tensor(
indices=non_contig_indices, values=weight.values(), size=weight.shape)
weight._coalesced_(True)
self.assertFalse(weight._indices().is_contiguous())
# Create un/coalesced sparse tensor
bias = torch.randn((hidden_size, 1), dtype=dtype, device=device).to_sparse()
bias = torch.cat([bias] * batch_size, dim=1)
if coalesced:
bias = bias.coalesce()
x = torch.randn(input_size, batch_size, dtype=dtype, device=device)
res = bias.sspaddmm(weight, x)
true_result = (bias.to_dense() + torch.matmul(weight.to_dense(), x)).to_sparse()
self.assertEqual(self.safeToDense(res), self.safeToDense(true_result))
@coalescedonoff
@precisionOverride({torch.bfloat16: 5e-2, torch.float16: 5e-2})
@dtypes(torch.double, torch.cdouble, torch.bfloat16, torch.float16)
def test_sparse_addmm(self, device, dtype, coalesced):
if (dtype is torch.bfloat16 or dtype is torch.float16) and device.startswith("cuda"):
self.skipTest('addmm_sparse_cuda is not implemented for BFloat16 and Half')
def test_shape(m, n, p, nnz, broadcast, alpha_beta=None):
if alpha_beta is None:
alpha = random.random()
beta = random.random()
else:
alpha, beta = alpha_beta
if broadcast:
D1 = make_tensor((), dtype=dtype, device=device, requires_grad=True)
else:
D1 = make_tensor([n, p], dtype=dtype, device=device, requires_grad=True)
D2 = make_tensor([m, p], dtype=dtype, device=device, requires_grad=True)
S = self._gen_sparse(2, nnz, [n, m], dtype, device, coalesced)[0]
S_dense = S.to_dense().requires_grad_(True)
S.requires_grad_(True)
Y = torch.sparse.addmm(D1, S, D2, beta=beta, alpha=alpha)
Y_dense = torch.addmm(D1, S_dense, D2, beta=beta, alpha=alpha)
self.assertEqual(Y, Y_dense)
if dtype not in {torch.double, torch.cdouble}:
# gradcheck will likely fail with low-precision input dtypes.
return
def fn(S, D1, D2, beta=beta, alpha=alpha):
return torch.sparse.addmm(D1, S, D2, beta=beta, alpha=alpha)
gradcheck(fn, (S, D1, D2), masked=True)
test_shape(7, 8, 9, 20, False, None)
test_shape(7, 8, 9, 20, True, None)
test_shape(7, 8, 9, 20, False, (1, 0))
test_shape(7, 8, 9, 20, True, (1, 0))
test_shape(7, 8, 9, 20, False, (1, 1))
test_shape(7, 8, 9, 20, True, (1, 1))
@coalescedonoff
@dtypes(torch.double)
@unittest.skipIf(TEST_WITH_CROSSREF, "generator unsupport triggers assertion error")
def test_sparse_mm(self, device, dtype, coalesced):
def test_shape(d1, d2, d3, nnz, transposed):
if transposed:
D = torch.randn(d3, d2, dtype=dtype,
device=device).t_().requires_grad_(True)
else:
D = torch.randn(d2, d3, dtype=dtype, device=device).requires_grad_(True)
S = self._gen_sparse(2, nnz, [d1, d2], dtype, device, coalesced)[0]
S_dense = S.to_dense().requires_grad_(True)
S.requires_grad_(True)
self.assertEqual(torch.sparse.mm(S, D), torch.mm(S_dense, D))
def fn(S, D):
return torch.sparse.mm(S, D)
gradcheck(fn, (S, D), masked=True)
test_shape(7, 8, 9, 20, False)
test_shape(7, 8, 9, 20, True)
@coalescedonoff
@dtypes(torch.double)
@unittest.skipIf(TEST_WITH_CROSSREF, "generator unsupport triggers assertion error")
@gradcheck_semantics()
def test_sparse_mul(self, device, dtype, coalesced, gradcheck):
# https://github.com/pytorch/pytorch/issues/79914
a = torch.tensor([[0., 1]], dtype=dtype, device=device).to_sparse().requires_grad_(True)
b = torch.tensor([[0., 1]], dtype=dtype, device=device).to_sparse().requires_grad_(True)
gradcheck(lambda x, y: torch.sparse.sum(x * y).to_dense(masked_grad=gradcheck.masked), [a, b])
def test_shape(sparse_dims, nnz, with_shape):
a = self._gen_sparse(sparse_dims, nnz, with_shape, dtype, device, coalesced)[0].requires_grad_(True)
b = self._gen_sparse(sparse_dims, nnz, with_shape, dtype, device, coalesced)[0].requires_grad_(True)
self.assertEqual((a * b).to_dense(), a.to_dense() * b.to_dense(), masked=True)
gradcheck(lambda x, y: (x * y).to_dense(), [a, b])
# Issues with 0-dim indices/values
gradcheck(lambda x, y: torch.sparse.sum(x * y).to_dense(), [a, b], masked=True)
# TODO: Re-enable these
# test_shape(2, 3, [2, 3, 4, 5])
# test_shape(2, 3, [2, 2, 0])
@coalescedonoff
@dtypes(torch.double)
def test_dsmm(self, device, dtype, coalesced):
def test_shape(di, dj, dk, nnz):
x = self._gen_sparse(2, nnz, [di, dj], dtype, device, coalesced)[0]
y = self.randn(dj, dk, dtype=dtype, device=device)
res = torch.dsmm(x, y)
expected = torch.mm(self.safeToDense(x), y)
self.assertEqual(res, expected)
test_shape(7, 5, 3, 20)
test_shape(1000, 100, 100, 20)
test_shape(3000, 64, 300, 20)
test_shape(0, 100, 100, 0)
test_shape(1000, 0, 100, 0)
test_shape(1000, 100, 0, 0)
test_shape(1000, 100, 0, 20)
@coalescedonoff
@dtypes(torch.double)
def test_hsmm(self, device, dtype, coalesced):
def test_shape(di, dj, dk, nnz):
x = self._gen_sparse(2, nnz, [di, dj], dtype, device, coalesced)[0]
y = self.randn(dj, dk, dtype=dtype, device=device)
res = torch.hsmm(x, y)
expected = torch.mm(self.safeToDense(x), y)
self.assertEqual(res.to_dense(), expected)
test_shape(7, 5, 3, 20)
test_shape(1000, 100, 100, 20)
test_shape(3000, 64, 300, 20)
test_shape(0, 100, 100, 0)
test_shape(1000, 0, 100, 0)
test_shape(1000, 100, 0, 0)
test_shape(1000, 100, 0, 20)
@coalescedonoff
@dtypes(torch.double)
def test_spadd(self, device, dtype, coalesced):
def _test_spadd_shape(nnz, shape_i, shape_v=None):
shape = shape_i + (shape_v or [])
x, _, _ = self._gen_sparse(len(shape_i), nnz, shape, dtype, device, coalesced)
y = self.randn(*shape, dtype=dtype, device=device)
r = random.random()
res = torch.add(y, x, alpha=r)
expected = y + r * self.safeToDense(x)
self.assertEqual(res, expected)
# Non contiguous dense tensor
s = list(shape)
s[0] = shape[-1]
s[-1] = shape[0]
y = self.randn(*s, dtype=dtype, device=device)
y.transpose_(0, len(s) - 1)
r = random.random()
res = torch.add(y, x, alpha=r)
expected = y + r * self.safeToDense(x)
self.assertEqual(res, expected)
x, i, v = self._gen_sparse(len(shape_i), nnz, shape, dtype, device, coalesced)
nnz = i.size(1)
# Non contiguous sparse indices tensor
x_ = self.sparse_tensor(i[:, ::2], v[:(nnz + 1) // 2], x.shape, dtype=dtype, device=device)
res = torch.add(y, x_, alpha=r)
expected = y + r * self.safeToDense(x_)
self.assertEqual(res, expected)
# Non contiguous sparse values tensor
x_ = self.sparse_tensor(i[:, :(nnz + 1) // 2], v[::2], x.shape, dtype=dtype, device=device)
res = torch.add(y, x_, alpha=r)
expected = y + r * self.safeToDense(x_)
self.assertEqual(res, expected)
# Non contiguous sparse indices and values tensors
x_ = self.sparse_tensor(i[:, 1::2], v[1::2], x.shape, dtype=dtype, device=device)
res = torch.add(y, x_, alpha=r)
expected = y + r * self.safeToDense(x_)
self.assertEqual(res, expected)
def _test_spadd():
_test_spadd_shape(10, [5, 6])
_test_spadd_shape(10, [10, 10, 10])
_test_spadd_shape(10, [50, 30, 20])
_test_spadd_shape(10, [5, 5, 5, 5, 5, 5])
_test_spadd_shape(0, [0, 30, 20])
_test_spadd_shape(0, [50, 0, 20])
_test_spadd_shape(0, [50, 30, 0])
def _test_spadd_hybrid():
_test_spadd_shape(10, [5, 6], [2, 3])
_test_spadd_shape(10, [10, 10, 10], [3])
_test_spadd_shape(10, [50, 30, 20], [2])
_test_spadd_shape(10, [5, 5, 5, 5, 5, 5], [2])
_test_spadd_shape(0, [0, 30, 20], [2, 0])
_test_spadd_shape(0, [50, 0, 20], [2, 0])
_test_spadd_shape(0, [50, 30, 0], [2, 0])
_test_spadd_shape(10, [50, 30, 20], [2, 0])
_test_spadd()
_test_spadd_hybrid()
@coalescedonoff
@dtypes(torch.float)
def test_sparse_add_out_bfloat16(self, device, dtype, coalesced):
# fp32
x, _, _ = self._gen_sparse(3, 5, 10, dtype, device, coalesced)
y, _, _ = self._gen_sparse(3, 5, 10, dtype, device, coalesced)
res_fp32 = torch.add(x, y)
# bfloat16
x = x.bfloat16()
y = y.bfloat16()
res_bf16 = torch.add(x, y)
res_bf16 = res_bf16.float() # to compare with reference
self.assertEqual(res_fp32, res_bf16, atol=1e-2, rtol=0)
@coalescedonoff
@dtypes(torch.double, torch.cdouble)
def test_norm(self, device, dtype, coalesced):
def test_shape(sparse_dims, nnz, with_size):
x, _, _ = self._gen_sparse(sparse_dims, nnz, with_size, dtype, device, coalesced)
y = x.coalesce()
self.assertEqual(x.norm(), y._values().norm())
test_shape(3, 10, 100)
test_shape(4, 10, [100, 100, 100, 5, 5, 5, 0])
test_shape(4, 0, [0, 0, 100, 5, 5, 5, 0])
# Unsupported arguments should error
kwarg_error_pairs = [
({'keepdim': True},
RuntimeError, r'norm_sparse currently does not support keepdim=True'),
({'dim': 0},
RuntimeError, r'norm_sparse currently only supports full reductions'),
({'dtype': torch.double, 'p': 'fro'},
ValueError, r'dtype argument is not supported in frobenius norm'),
({'dtype': torch.double, 'p': 0},
RuntimeError, r"norm_sparse currently does not support 'dtype' argument")
]
x = self._gen_sparse(3, 10, 100, dtype, device, coalesced)[0]
for kwargs, err, msg in kwarg_error_pairs:
with self.assertRaisesRegex(err, msg):
x.norm(**kwargs)
@coalescedonoff
@dtypes(torch.double)
@unittest.skipIf(TEST_WITH_CROSSREF, "fallback triggers cuda device error")
def test_sparse_sum(self, device, dtype, coalesced):
def run_tests(S, td=None):
D = S.coalesce().to_dense().detach().requires_grad_(True)
if td is None:
S_sum = torch.sparse.sum(S)
D_sum = D.sum()
self.assertEqual(S_sum.item(), D_sum.item())
def fn(S):
return torch.sparse.sum(S)
gradcheck(fn, (S,), masked=True)
else:
S_sum = torch.sparse.sum(S, td)
D_sum = D.sum(td)
self.assertEqual(S_sum.to_dense() if S_sum.is_sparse else S_sum, D_sum)
def fn(S):
res = torch.sparse.sum(S, td)
return res.to_dense(masked_grad=True)
gradcheck(fn, (S,), masked=True)
nnz = 10
sparse_dims = 2
with_size = [5, 5, 1, 4] # use a dense dim = 1 to test for squeeze
test_dims = []
for i in range(1, 5):
test_dims += itertools.combinations(range(len(with_size)), i)
# https://github.com/pytorch/pytorch/issues/16501
x = torch.tensor([[1., 0., 0., 1.],
[0., 1., 0., 0.],
[0., 1., 1., 0.],
[0., 1., 0., 2.]], dtype=dtype, device=device).to_sparse()
self.assertEqual(torch.sparse.sum(x, dim=0), torch.sparse.sum(x, dim=-2))
self.assertEqual(torch.sum(x.to_dense(), dim=0), torch.sparse.sum(x, dim=0).to_dense())
S = self._gen_sparse(sparse_dims, nnz, with_size, dtype, device, coalesced)[0]
# dim out of range
self.assertRaises(IndexError, lambda: torch.sparse.sum(S, 5))
# dim 0 appears multiple times in the list of dims
self.assertRaises(RuntimeError, lambda: torch.sparse.sum(S, [0, 0]))
# sum an empty tensor
empty_S = torch.sparse_coo_tensor(size=with_size, dtype=dtype, device=device)
self.assertEqual(torch.sparse.sum(empty_S, [0]).to_dense(), torch.sum(empty_S.to_dense(), [0]))
self.assertEqual(torch.sparse.sum(empty_S), torch.tensor(0, dtype=dtype, device=device))
empty_S.requires_grad_(True)
empty_S_sum = torch.sparse.sum(empty_S)
empty_S_sum.backward()
self.assertEqual(empty_S.grad.to_dense(), empty_S.clone().detach().to_dense())
# test values().sum()
S = self._gen_sparse(sparse_dims, nnz, with_size, dtype, device, coalesced)[0]
run_tests(S.requires_grad_(True))
for test_dim in test_dims:
S = self._gen_sparse(sparse_dims, nnz, with_size, dtype, device, coalesced)[0]
run_tests(S.requires_grad_(True), test_dim)
def _test_basic_ops_shape(self, nnz_x1, nnz_x2, shape_i, shape_v, dtype, device, coalesced):
shape = shape_i + (shape_v)
x1, _, _ = self._gen_sparse(len(shape_i), nnz_x1, shape, dtype, device, coalesced)
x2, _, _ = self._gen_sparse(len(shape_i), nnz_x2, shape, dtype, device, coalesced)
y1 = x1 + x2
y2 = x1.clone()
y2.add_(x2)
expected = self.safeToDense(x1) + self.safeToDense(x2)
self.assertEqual(self.safeToDense(y1), expected)
self.assertEqual(self.safeToDense(y2), expected)
y1 = x1 - x2
y2 = x1.clone()
y2.sub_(x2)
expected = self.safeToDense(x1) - self.safeToDense(x2)
self.assertEqual(self.safeToDense(y1), expected)
self.assertEqual(self.safeToDense(y2), expected)
y1 = x1 * x2
y2 = x1.clone()
y2.mul_(x2)
expected = self.safeToDense(x1) * self.safeToDense(x2)
self.assertEqual(self.safeToDense(y1), expected)
self.assertEqual(self.safeToDense(y2), expected)
y1 = x1 * 37.5
y2 = x1.clone()
y2.mul_(37.5)
expected = self.safeToDense(x1) * 37.5
self.assertEqual(self.safeToDense(y1), expected)
self.assertEqual(self.safeToDense(y2), expected)
y1 = x1 / 37.5
y2 = x1.clone()
y2.div_(37.5)
expected = self.safeToDense(x1) / 37.5
self.assertEqual(self.safeToDense(y1), expected)
self.assertEqual(self.safeToDense(y2), expected)
y1 = x1 // 37.5
y2 = x1.clone()
y2.floor_divide_(37.5)
expected = self.safeToDense(x1) // 37.5
self.assertEqual(self.safeToDense(y1), expected)
self.assertEqual(self.safeToDense(y2), expected)
# TODO: add back inplace support
y1 = x1 ** 2
y2 = x1.clone()
y2 = y2.pow(2)
expected = self.safeToDense(x1) ** 2
self.assertEqual(self.safeToDense(y1), expected)
self.assertEqual(self.safeToDense(y2), expected)
y = x1.clone()
y.zero_()
expected = torch.zeros(x1.size(), dtype=dtype, device=device)
self.assertEqual(self.safeToDense(y), expected)
self.assertEqual(x1.is_coalesced(), coalesced)
y = x1.coalesce()
z = x1.coalesce()
self.assertEqual(x1.is_coalesced(), coalesced)
self.assertTrue(y.is_coalesced())
y._values().add_(1)
if not x1.is_coalesced():
# check that coalesce is out of place if the original tensor is not
# coalesced.
self.assertEqual(z._values() + 1, y._values())
else:
# check that coalesce is in-place if the original tensor is
# coalesced.
self.assertEqual(z._values(), y._values())
@coalescedonoff
@dtypes(torch.double)
def test_basic_ops(self, device, dtype, coalesced):
def _test_basic_ops():
self._test_basic_ops_shape(9, 12, [5, 6], [], dtype, device, coalesced)
self._test_basic_ops_shape(9, 12, [10, 10, 10], [], dtype, device, coalesced)
self._test_basic_ops_shape(9, 12, [50, 30, 20], [], dtype, device, coalesced)
self._test_basic_ops_shape(9, 12, [5, 5, 5, 5, 5, 5], [], dtype, device, coalesced)
self._test_basic_ops_shape(0, 12, [10, 10, 10], [], dtype, device, coalesced)
self._test_basic_ops_shape(9, 0, [10, 10, 10], [], dtype, device, coalesced)
self._test_basic_ops_shape(0, 0, [10, 10, 10], [], dtype, device, coalesced)
self._test_basic_ops_shape(0, 0, [10, 10, 0], [], dtype, device, coalesced)
self._test_basic_ops_shape(0, 0, [], [], dtype, device, coalesced)
def _test_basic_ops_hybrid():
self._test_basic_ops_shape(9, 12, [5, 6], [2, 3], dtype, device, coalesced)
self._test_basic_ops_shape(9, 12, [10, 10, 10], [3], dtype, device, coalesced)
self._test_basic_ops_shape(9, 12, [50, 30, 20], [2], dtype, device, coalesced)
self._test_basic_ops_shape(9, 12, [5, 5, 5, 5, 5, 5], [2], dtype, device, coalesced)
self._test_basic_ops_shape(0, 12, [10, 10, 10], [2], dtype, device, coalesced)
self._test_basic_ops_shape(9, 0, [10, 10, 10], [2], dtype, device, coalesced)
self._test_basic_ops_shape(0, 0, [10, 10, 10], [2], dtype, device, coalesced)
self._test_basic_ops_shape(9, 12, [10, 10, 10], [2, 0], dtype, device, coalesced)
self._test_basic_ops_shape(0, 12, [10, 10, 10], [2, 0], dtype, device, coalesced)
self._test_basic_ops_shape(9, 0, [10, 10, 10], [2, 0], dtype, device, coalesced)
self._test_basic_ops_shape(0, 0, [10, 10, 10], [2, 0], dtype, device, coalesced)
self._test_basic_ops_shape(0, 0, [10, 10, 0], [2, 0], dtype, device, coalesced)
_test_basic_ops()
_test_basic_ops_hybrid()
@dtypes(torch.double, torch.cdouble)
def test_add_dense_sparse_mismatch(self, device, dtype):
def test_shape(dense_size, sparse_dims_shape, dense_dims_shape, sparse_size):
x = torch.zeros(dense_size, dtype=dtype, device=device)
sparse_y = self.sparse_tensor(torch.zeros(sparse_dims_shape, dtype=torch.int64, device=device),
torch.randn(dense_dims_shape, dtype=dtype, device=device),
torch.Size(sparse_size))
with self.assertRaisesRegex(
RuntimeError,
"add: expected 'self' and 'other' to have same size"):
x + sparse_y
test_shape([3, 4], [1, 4], [4, 4, 4], [3, 4, 4])
test_shape([3, 4, 0], [1, 4], [4, 4, 4, 0], [3, 4, 4, 0])
@skipIfTorchDynamo("Not a TorchDynamo suitable test")
@dtypes(torch.double, torch.cdouble)
def test_add_noncontiguous(self, device, dtype):
indices = self.index_tensor([[1, 2], [0, 2]], device=device)
values = torch.tensor([1.], dtype=dtype, device=device).expand(2, 3, 4, 5)
x = self.sparse_tensor(indices, values, dtype=dtype, device=device)
assert not x._values().is_contiguous()
y = x + x
expected = self.safeToDense(x) + self.safeToDense(x)
self.assertEqual(self.safeToDense(y), expected)
def _test_sparse_mask_shape(self, nnz_x1, nnz_x2, shape_i, shape_v, dtype, device, coalesced):
shape = shape_i + (shape_v or [])
x1, _, _ = self._gen_sparse(len(shape_i), nnz_x1, shape, dtype, device, coalesced)
x2, _, _ = self._gen_sparse(len(shape_i), nnz_x2, shape, dtype, device, coalesced)
y1 = x1 + x2
y2 = x1.clone()
y2.add_(x2)
expected = self.safeToDense(x1) + self.safeToDense(x2)
self.assertEqual(self.safeToDense(y1), expected)
self.assertEqual(self.safeToDense(y2), expected)
@coalescedonoff
@dtypes(torch.double, torch.cdouble)
def test_sparse_mask(self, device, dtype, coalesced):
def _test_sparse_mask_fixed():
i = self.index_tensor([
[1, 3, 0, 4],
[2, 1, 2, 3],
], device=device)
v = torch.tensor([1, 2, 3, 4], dtype=dtype, device=device)
x = self.sparse_tensor(i, v, torch.Size([5, 4]), dtype=dtype, device=device).coalesce()
dense = torch.tensor([
[1, 2, 3, 4],
[5, 6, 7, 8],
[9, 10, 11, 12],
[13, 14, 15, 16],
[17, 18, 19, 20],
], dtype=dtype, device=device)
exp_v = torch.tensor([7, 14, 3, 20], dtype=dtype, device=device)
res_dense_lhs = dense.sparse_mask(x)
sparse = dense.to_sparse()
res_sparse_lhs = sparse.sparse_mask(x)
expected = self.sparse_tensor(i, exp_v, torch.Size([5, 4]), dtype=dtype, device=device)
self.assertEqual(res_dense_lhs.coalesce(), expected.coalesce())
# check no side effects for the coalesce flag.
self.assertTrue(sparse.is_coalesced())
self.assertEqual(res_sparse_lhs.coalesce(), expected.coalesce())
i = self.index_tensor([
[1, 3, 0, 4],
[2, 1, 2, 3],
], device=device)
v = torch.empty([4, 0], dtype=dtype, device=device)
x = self.sparse_tensor(i, v, torch.Size([5, 4, 0])).coalesce()
dense = torch.empty([5, 4, 0], dtype=dtype, device=device)
exp_v = torch.empty([4, 0], dtype=dtype, device=device)
res_dense_lhs = dense.sparse_mask(x)
sparse = dense.to_sparse(2)
res_sparse_lhs = sparse.sparse_mask(x)
expected = self.sparse_tensor(i, exp_v, torch.Size([5, 4, 0]), dtype=dtype, device=device)
self.assertEqual(res_dense_lhs.coalesce(), expected.coalesce())
# check no side effects for the coalesce flag.
self.assertTrue(sparse.is_coalesced())
self.assertEqual(res_sparse_lhs.coalesce(), expected.coalesce())
_test_sparse_mask_fixed()
self._test_sparse_mask_shape(9, 12, [5, 6], [], dtype, device, coalesced)
self._test_sparse_mask_shape(9, 12, [10, 10, 10], [], dtype, device, coalesced)
self._test_sparse_mask_shape(9, 12, [50, 30, 20], [], dtype, device, coalesced)
self._test_sparse_mask_shape(9, 12, [5, 5, 5, 5, 5, 5], [], dtype, device, coalesced)
self._test_sparse_mask_shape(0, 12, [10, 10, 10], [], dtype, device, coalesced)
self._test_sparse_mask_shape(9, 0, [10, 10, 10], [], dtype, device, coalesced)
self._test_sparse_mask_shape(0, 0, [10, 10, 10], [], dtype, device, coalesced)
self._test_sparse_mask_shape(0, 0, [10, 10, 0], [], dtype, device, coalesced)
# check repetitions and matchings in the intersection
lhs = torch.randint(0, 5, (100,), device=device)
rhs = torch.randint(0, 5, (100,), device=device).to_sparse()
self.assertEqual(lhs.to_sparse().sparse_mask(rhs), lhs.sparse_mask(rhs))
# check coalesce
sparse_c = torch.rand(3, 3, device=device).to_sparse()
sparse_unc = torch.rand(3, 3, device=device).to_sparse()._coalesced_(False)
for lhs, rhs in [(sparse_c, sparse_unc), (sparse_unc, sparse_c)]:
res_all_sparse = lhs.sparse_mask(rhs)
res_dense_sparse = lhs.to_dense().sparse_mask(rhs)
self.assertEqual(res_all_sparse.coalesce(), res_dense_sparse.coalesce())
self.assertEqual(rhs.is_coalesced(), res_all_sparse.is_coalesced())
@coalescedonoff
@dtypes(torch.double, torch.cdouble)
def test_sparse_mask_hybrid(self, device, dtype, coalesced):
def _test_sparse_mask_hybrid_fixed():
i = self.index_tensor([
[1, 3, 0, 4],
[2, 1, 2, 3],
])
v = torch.tensor([[1, 2], [2, 3], [3, 4], [4, 5]])
# TODO: This is also testing that, if coalesce is a no-op,
# the indices don't get permuted. I don't know if we actually
# want to give this invariant.
x = self.sparse_tensor(i, v, torch.Size([5, 4, 2])).coalesce()
dense = torch.tensor([
[[1, 3], [2, 2], [3, 3], [4, 2]],
[[5, 7], [6, 7], [7, 9], [8, 9]],
[[9, 2], [10, 4], [11, 1], [12, 3]],
[[13, 5], [14, 1], [15, 1], [16, 6]],
[[17, 7], [18, 2], [19, 7], [20, 1]],
])
res_dense_lhs = dense.sparse_mask(x)
sparse = dense.to_sparse(2)
res_sparse_lhs = sparse.sparse_mask(x)
exp_v = torch.tensor([[7, 9], [14, 1], [3, 3], [20, 1]])
expected = self.sparse_tensor(i, exp_v, torch.Size([5, 4, 2]))
self.assertEqual(res_dense_lhs.coalesce(), expected.coalesce())
# check no side effects for the coalesce flag
self.assertTrue(sparse.is_coalesced())
self.assertEqual(res_sparse_lhs.coalesce(), expected.coalesce())
i = self.index_tensor([
[1, 3, 0, 4],
[2, 1, 2, 3],
])
v = torch.empty(4, 2, 0)
x = self.sparse_tensor(i, v, torch.Size([5, 4, 2, 0])).coalesce()
dense = torch.empty(5, 4, 2, 0)
res_dense_lhs = dense.sparse_mask(x)
sparse = dense.to_sparse(2)
res_sparse_lhs = sparse.sparse_mask(x)
exp_v = torch.empty(4, 2, 0)
expected = self.sparse_tensor(i, exp_v, torch.Size([5, 4, 2, 0]))
self.assertEqual(res_dense_lhs.coalesce(), expected.coalesce())
# check no side effects for the coalesce flag
self.assertTrue(sparse.is_coalesced())
self.assertEqual(res_sparse_lhs.coalesce(), expected.coalesce())
_test_sparse_mask_hybrid_fixed()
self._test_sparse_mask_shape(9, 12, [5, 6], [2, 3], dtype, device, coalesced)
self._test_sparse_mask_shape(9, 12, [10, 10, 10], [3], dtype, device, coalesced)
self._test_sparse_mask_shape(9, 12, [50, 30, 20], [2], dtype, device, coalesced)
self._test_sparse_mask_shape(9, 12, [5, 5, 5, 5, 5, 5], [2], dtype, device, coalesced)
self._test_sparse_mask_shape(0, 12, [10, 10, 10], [2], dtype, device, coalesced)
self._test_sparse_mask_shape(9, 0, [10, 10, 10], [2], dtype, device, coalesced)
self._test_sparse_mask_shape(0, 0, [10, 10, 10], [2], dtype, device, coalesced)
self._test_sparse_mask_shape(9, 12, [10, 10, 10], [2, 0], dtype, device, coalesced)
self._test_sparse_mask_shape(0, 12, [10, 10, 10], [2, 0], dtype, device, coalesced)
self._test_sparse_mask_shape(9, 0, [10, 10, 10], [2, 0], dtype, device, coalesced)
self._test_sparse_mask_shape(0, 0, [10, 10, 10], [2, 0], dtype, device, coalesced)
self._test_sparse_mask_shape(0, 0, [10, 10, 0], [2, 0], dtype, device, coalesced)
@dtypes(torch.double, torch.cdouble)
@skipIfCrossRef
def test_sparse_mask_backward(self, device, dtype):
from itertools import product, repeat
shape = (5, 5)
sparse_dims = len(shape)
nnzs = (0, 5, 15, 25)
lhs_data = torch.arange(1, 26, device=device).reshape(shape).to(dtype).to_sparse(sparse_dims)
rhs_data = lhs_data.clone()
for nnz in nnzs:
for lhs_is_coalesced, rhs_is_coalesced in product(*repeat((True, False), 2)):
lhs = torch.sparse_coo_tensor(
lhs_data._indices()[:, :nnz],
lhs_data._values()[:nnz],
lhs_data.shape
).clone()._coalesced_(lhs_is_coalesced).requires_grad_(True)
rhs = torch.sparse_coo_tensor(
lhs_data._indices()[:, -nnz:],
lhs_data._values()[-nnz:],
lhs_data.shape
).clone()._coalesced_(rhs_is_coalesced)
# To test masked semantics we need to make sure that
# sparsity_pattern(lhs) == sparsity_pattern(lhs.grad).
# lhs.sparse_mask(lhs_mask) accomplishes that.
lhs_mask = lhs.detach().clone()
gradcheck(lambda x: x.sparse_mask(lhs_mask).sparse_mask(rhs).to_dense(masked_grad=True), (lhs,), masked=True)
gradcheck(lambda x: x.sparse_mask(rhs).to_dense(masked_grad=False), (lhs,), masked=False)
@coalescedonoff
@dtypes(torch.double, torch.cdouble)
def test_zeros(self, device, dtype, coalesced):
def _test_zeros(nnzs, shape, out_shape_i, out_shape_v=None):
out_shape = out_shape_i + (out_shape_v or [])
for nnz in nnzs:
out, _, _ = self._gen_sparse(len(out_shape_i), nnz, out_shape, dtype, device, coalesced)
torch.zeros(*shape, out=out, dtype=dtype, device=device)
self.assertEqual(tuple(out.size()), tuple(shape))
self.assertTrue(out._indices().numel() == out._values().numel() == 0)
self.assertEqual(out._nnz(), 0)
self.assertEqual(out.sparse_dim(), len(shape))
self.assertEqual(out.dense_dim(), 0)
def test_shape(i_shapes, v_shapes, shape, nnzs):
for i_dim in range(1, len(i_shapes) + 1):
for v_dim in range(len(v_shapes) + 1):
_test_zeros(nnzs, shape, i_shapes[:i_dim], v_shapes[:v_dim])
test_shape([2, 3, 4], [3, 4, 5, 6], [2, 3, 4], [9, 12])
test_shape([0, 3, 4], [3, 4, 5, 6], [2, 3, 4], [0])
test_shape([2, 3, 4], [0, 4, 5, 6], [2, 3, 4], [9, 12])
test_shape([2, 3, 4], [3, 4, 5, 6], [2, 3, 0], [9, 12])
test_shape([0, 3, 4], [3, 4, 5, 6], [2, 3, 0], [0])
test_shape([2, 3, 4], [0, 4, 5, 6], [2, 3, 0], [9, 12])
@coalescedonoff
@dtypes(torch.double, torch.cdouble)
def test_zeros_like(self, device, dtype, coalesced):
def _test_zeros_like(nnzs, template_shape_i, template_shape_v=None):
template_shape_v = template_shape_v or []
template_shape = template_shape_i + template_shape_v
for nnz in nnzs:
t, _, _ = self._gen_sparse(len(template_shape_i), nnz, template_shape, dtype, device, coalesced)
res = torch.zeros_like(t)
self.assertEqual(tuple(res.size()), tuple(template_shape))
self.assertTrue(res._indices().numel() == res._values().numel() == 0)
self.assertEqual(res._nnz(), 0)
self.assertEqual(res.sparse_dim(), len(template_shape_i))
self.assertEqual(res.dense_dim(), len(template_shape_v))
def test_shape(i_shapes, v_shapes, nnzs):
for i_dim in range(1, len(i_shapes) + 1):
for v_dim in range(len(v_shapes) + 1):
_test_zeros_like(nnzs, i_shapes[:i_dim], v_shapes[:v_dim])
test_shape([2, 3, 4], [3, 4, 5, 6], [9, 12])
test_shape([0, 3, 4], [3, 4, 5, 6], [0])
test_shape([2, 3, 4], [0, 4, 5, 6], [9, 12])
test_shape([2, 3, 4], [3, 4, 5, 6], [9, 12])
test_shape([0, 3, 4], [3, 4, 5, 6], [0])
test_shape([2, 3, 4], [0, 4, 5, 6], [9, 12])
sparse_tensor, _, _ = self._gen_sparse(len([2, 3]), 9, [2, 3] + [5, 6], dtype, device, coalesced)
data = (sparse_tensor, sparse_tensor, sparse_tensor, sparse_tensor.unsqueeze(0))
mem_formats = [torch.channels_last, torch.contiguous_format, torch.preserve_format, torch.channels_last_3d]
for x, mem_format in zip(data, mem_formats):
with self.assertRaisesRegex(RuntimeError, "memory format option is only supported by strided tensors"):
result = torch.zeros_like(x, memory_format=mem_format)
result = torch.zeros_like(x, layout=torch.strided, memory_format=mem_format)
self.assertTrue(result.layout == torch.strided)
dense_tensor = sparse_tensor.to_dense()
result = torch.zeros_like(dense_tensor, layout=torch.sparse_coo)
self.assertEqual(dense_tensor.shape, result.shape)
self.assertEqual(result.layout, torch.sparse_coo)
sparse_zeros = torch.sparse_coo_tensor(dense_tensor.shape)
self.assertEqual(result._indices().shape, sparse_zeros._indices().shape)
self.assertEqual(result._values().shape, sparse_zeros._values().shape)
def _assert_sparse_invars(self, t):
# SparseTensor has the following invariants:
# - sparse_dim + dense_dim = len(SparseTensor.shape)
# - SparseTensor._indices().shape = (sparse_dim, nnz)
# - SparseTensor._values().shape = (nnz, SparseTensor.shape[sparse_dim:])
self.assertEqual(t.sparse_dim() + t.dense_dim(), len(t.shape))
self.assertEqual(tuple(t._indices().shape), (t.sparse_dim(), t._nnz()))
self.assertEqual(tuple(t._values().shape), (t._nnz(), ) + t.shape[t.sparse_dim():])
def _test_empty_like(self, sparse_tensor, dtype, device, coalesced):
result = torch.empty_like(sparse_tensor)
self.assertTrue(result.is_sparse)
self._assert_sparse_invars(result)
self.assertEqual(result.shape, sparse_tensor.shape)
self.assertEqual(result.dtype, sparse_tensor.dtype)
self.assertEqual(result.device, sparse_tensor.device)
self.assertEqual(result.sparse_dim(), sparse_tensor.sparse_dim())
self.assertEqual(result.dense_dim(), sparse_tensor.dense_dim())
sparse_tensor, _, _ = self._gen_sparse(len([2, 3]), 9, [2, 3] + [5, 6], dtype, device, coalesced)
data = (sparse_tensor, sparse_tensor, sparse_tensor, sparse_tensor.unsqueeze(0))
mem_formats = [torch.channels_last, torch.contiguous_format, torch.preserve_format, torch.channels_last_3d]
for x, mem_format in zip(data, mem_formats):
with self.assertRaisesRegex(RuntimeError, "memory format option is only supported by strided tensors"):
result = torch.empty_like(x, memory_format=mem_format)
result = torch.empty_like(x, layout=torch.strided, memory_format=mem_format)
self.assertTrue(result.layout == torch.strided)
with self.assertRaisesRegex(
RuntimeError, r"Could not run 'aten::empty_strided' with arguments from the 'Sparse(CPU|CUDA)' backend"
):
dense_tensor = sparse_tensor.to_dense()
result = torch.empty_like(dense_tensor, layout=torch.sparse_coo)
@coalescedonoff
@dtypes(torch.double, torch.cdouble)
def test_empty_like(self, device, dtype, coalesced):
# tests https://github.com/pytorch/pytorch/issues/43699
if coalesced:
input_coalesced = torch.sparse_coo_tensor(
indices=torch.tensor([[0, 1, 2]]),
values=torch.tensor([3.0, -4.0, 5.0]),
size=[3, ],
dtype=dtype,
device=device
).coalesce()
self._test_empty_like(input_coalesced, dtype, device, coalesced)
# hybrid sparse input
input_coalesced = torch.sparse_coo_tensor(
indices=torch.tensor([[1, 3], [2, 4]]),
values=torch.tensor([[-1.0, 3.0], [-5.0, 7.0]]),
size=[4, 5, 2],
dtype=dtype,
device=device
).coalesce()
self._test_empty_like(input_coalesced, dtype, device, coalesced)
if not coalesced:
# test uncoalesced input
input_uncoalesced = torch.sparse_coo_tensor(
indices=torch.tensor([[0], [1], [2], [0], [1], [2]]).transpose(1, 0),
values=torch.tensor([2.0, -3.0, -4.0, 1.0, -1.0, 1.5]),
size=[3, ],
dtype=dtype,
device=device
)
self._test_empty_like(input_uncoalesced, dtype, device, coalesced)
# test on empty sparse tensor
input_uncoalesced = torch.sparse_coo_tensor(
indices=torch.zeros([2, 0]),
values=torch.zeros([0, 5, 5, 5, 5, 5, 5, 0]),
size=[0, 0, 5, 5, 5, 5, 5, 5, 0],
dtype=dtype,
device=device
)
self._test_empty_like(input_uncoalesced, dtype, device, coalesced)
def _test_narrow(self, input, narrow_args):
expected = input.to_dense().narrow(*narrow_args)
self.assertEqual(expected, input.narrow_copy(*narrow_args).to_dense())
def _all_narrow_combs(self, shape):
for dim, dim_sz in enumerate(shape):
for start in range(dim_sz):
for length in range(dim_sz - start):
yield [dim, start, length]
@coalescedonoff
@dtypes(torch.double, torch.cdouble)
def test_narrow(self, device, dtype, coalesced):
shape = [3, 3, 4, 2]
input, _, _ = self._gen_sparse(4, 19, shape, dtype, device, coalesced)
for narrow_args in self._all_narrow_combs(shape):
self._test_narrow(input, narrow_args)
self.assertRaises(RuntimeError, lambda: input.narrow_copy(-1, 0, 3)) # dim < 0
self.assertRaises(RuntimeError, lambda: input.narrow_copy(10, 0, 3)) # dim > input.dim()
self.assertRaises(RuntimeError, lambda: input.narrow_copy(0, shape[0] + 1, 3)) # start > size of dim
self.assertRaises(RuntimeError, lambda: input.narrow_copy(0, 2, shape[0])) # start+length > size of dim
with_dense, _, _ = self._gen_sparse(2, 7, shape, dtype, device, coalesced)
for narrow_args in self._all_narrow_combs(shape):
self._test_narrow(with_dense, narrow_args)
self.assertRaises(RuntimeError, lambda: with_dense.narrow_copy(10, 0, 3)) # dim > sparseDim + denseDim
def _test_log1p_tensor(self, sparse_tensor, coalesced):
def is_integral(dtype):
return dtype in integral_types()
dense_tensor = sparse_tensor.to_dense()
expected_output = dense_tensor.log1p()
is_integral_dtype = is_integral(sparse_tensor.dtype)
self.assertEqual(expected_output, sparse_tensor.log1p().to_dense())
if is_integral_dtype:
with self.assertRaisesRegex(RuntimeError, "result type .* can't be cast to"):
sparse_tensor.coalesce().log1p_()
else:
self.assertEqual(expected_output, sparse_tensor.coalesce().log1p_().to_dense())
if not coalesced:
# test in-place op on uncoalesced input
with self.assertRaisesRegex(RuntimeError, "log1p_ requires coalesced input"):
sparse_tensor.log1p_()
if is_integral_dtype:
with self.assertRaisesRegex(RuntimeError, "only Tensors of floating point dtype can require gradients"):
sparse_tensor.requires_grad_()
@coalescedonoff
@dtypes(*all_types())
def test_log1p(self, device, dtype, coalesced):
if coalesced:
input_coalesced = torch.sparse_coo_tensor(
indices=torch.tensor([[0], [1], [2]]).transpose(1, 0),
values=torch.tensor([3.0, 4.0, 5.0]),
size=[3, ],
device=device,
dtype=dtype
).coalesce()
self._test_log1p_tensor(input_coalesced, coalesced)
# hybrid sparse input
input_coalesced = torch.sparse_coo_tensor(
indices=torch.tensor([[1, 3], [2, 4]]),
values=torch.tensor([[1.0, 3.0], [5.0, 7.0]]),
size=[4, 5, 2],
device=device,
dtype=dtype
).coalesce()
self._test_log1p_tensor(input_coalesced, coalesced)
if not coalesced:
# test uncoalesced input
input_uncoalesced = torch.sparse_coo_tensor(
indices=torch.tensor([[0], [1], [2], [0], [1], [2]]).transpose(1, 0),
values=torch.tensor([2.0, 3.0, 4.0, 1.0, 1.0, 1.0]),
size=[3, ],
device=device,
dtype=dtype
)
self._test_log1p_tensor(input_uncoalesced, coalesced)
# test on empty sparse tensor
input_uncoalesced = torch.sparse_coo_tensor(
indices=torch.zeros([2, 0]),
values=torch.zeros([0, 5, 5, 5, 5, 5, 5, 0]),
size=[0, 0, 5, 5, 5, 5, 5, 5, 0],
device=device,
dtype=dtype
)
# empty tensors are coalesced at creation (nnz < 2) we must force the uncoalesced state
input_uncoalesced._coalesced_(False)
self._test_log1p_tensor(input_uncoalesced, coalesced)
def _test_neg_negative(self, sparse_tensor):
dense_tensor = sparse_tensor.to_dense()
expected_output = dense_tensor.neg()
ops = (
torch.neg, torch.Tensor.neg, torch.Tensor.neg_,
torch.negative, torch.Tensor.negative, torch.Tensor.negative_,
operator.neg
)
for op in ops:
sparse_tensor_copy = sparse_tensor.clone()
self.assertEqual(expected_output, op(sparse_tensor_copy).to_dense())
if op in (torch.neg, torch.negative):
sparse_tensor_out = torch.zeros_like(sparse_tensor)
op(sparse_tensor, out=sparse_tensor_out)
self.assertEqual(expected_output, sparse_tensor_out.to_dense())
@coalescedonoff
@dtypes(torch.double, torch.cdouble)
def test_neg_negative(self, device, dtype, coalesced):
if coalesced:
input_coalesced = torch.sparse_coo_tensor(
indices=torch.tensor([[0, 1, 2]]),
values=torch.tensor([3.0, -4.0, 5.0]),
size=[3, ],
dtype=dtype,
device=device
).coalesce()
self._test_neg_negative(input_coalesced)
# hybrid sparse input
input_coalesced = torch.sparse_coo_tensor(
indices=torch.tensor([[1, 3], [2, 4]]),
values=torch.tensor([[-1.0, 3.0], [-5.0, 7.0]]),
size=[4, 5, 2],
dtype=dtype,
device=device
).coalesce()
self._test_neg_negative(input_coalesced)
if not coalesced:
# test uncoalesced input
input_uncoalesced = torch.sparse_coo_tensor(
indices=torch.tensor([[0], [1], [2], [0], [1], [2]]).transpose(1, 0),
values=torch.tensor([2.0, -3.0, -4.0, 1.0, -1.0, 1.5]),
size=[3, ],
dtype=dtype,
device=device
)
self._test_neg_negative(input_uncoalesced)
# test on empty sparse tensor
input_uncoalesced = torch.sparse_coo_tensor(
indices=torch.zeros([2, 0]),
values=torch.zeros([0, 5, 5, 5, 5, 5, 5, 0]),
size=[0, 0, 5, 5, 5, 5, 5, 5, 0],
dtype=dtype,
device=device
)
self._test_neg_negative(input_uncoalesced)
def _test_asin_arcsin(self, sparse_tensor, coalesced):
def is_integral(dtype):
return dtype in integral_types()
is_integral_dtype = is_integral(sparse_tensor.dtype)
dense_tensor = sparse_tensor.to_dense()
expected_output = dense_tensor.asin()
ops = (
torch.asin, torch.Tensor.asin,
torch.arcsin, torch.Tensor.arcsin,
)
for op in ops:
self.assertEqual(expected_output, op(sparse_tensor).to_dense())
if op in (torch.asin, torch.arcsin):
sparse_tensor_out = torch.zeros_like(sparse_tensor)
if not is_integral_dtype:
op(sparse_tensor, out=sparse_tensor_out)
self.assertEqual(expected_output, sparse_tensor_out.to_dense())
else:
with self.assertRaisesRegex(RuntimeError, "result type .* can't be cast to"):
op(sparse_tensor, out=sparse_tensor_out)
for op in (torch.Tensor.asin_, torch.Tensor.arcsin_):
if is_integral_dtype:
# test coalesce on integral dtype tensor
with self.assertRaisesRegex(RuntimeError, "result type .* can't be cast to"):
op(sparse_tensor.clone().coalesce()).to_dense()
else:
self.assertEqual(expected_output, op(sparse_tensor.clone().coalesce()).to_dense())
if not coalesced:
# test in-place op on uncoalesced input
with self.assertRaisesRegex(RuntimeError, "asin_ requires coalesced input"):
op(sparse_tensor)
@coalescedonoff
@dtypes(*all_types())
def test_asin_arcsin(self, device, dtype, coalesced):
if coalesced:
input_coalesced = torch.sparse_coo_tensor(
indices=torch.tensor([[0, 1, 2, 3]]),
values=torch.tensor([0.5, -0.5, 0.7, -0.7]),
size=[4, ],
dtype=dtype,
device=device
).coalesce()
self._test_asin_arcsin(input_coalesced, coalesced)
# hybrid sparse input
input_coalesced = torch.sparse_coo_tensor(
indices=torch.tensor([[1, 3], [2, 4]]),
values=torch.tensor([[-0.1, 0.24], [-0.44, 0.1]]),
size=[4, 5, 2],
dtype=dtype,
device=device
).coalesce()
self._test_asin_arcsin(input_coalesced, coalesced)
if not coalesced:
# test uncoalesced input
input_uncoalesced = torch.sparse_coo_tensor(
indices=torch.tensor([[0], [1], [2], [0], [1], [2]]).transpose(1, 0),
values=torch.tensor([0.3, -0.3, -0.4, 0.3, -0.5, 0.15]),
size=[3, ],
dtype=dtype,
device=device
)
self._test_asin_arcsin(input_uncoalesced, coalesced)
# test on empty sparse tensor
input_uncoalesced = torch.sparse_coo_tensor(
indices=torch.zeros([2, 0]),
values=torch.zeros([0, 5, 5, 5, 5, 5, 5, 0]),
size=[0, 0, 5, 5, 5, 5, 5, 5, 0],
dtype=dtype,
device=device
)
# empty tensors are coalesced at creation (nnz < 2) we must force the uncoalesced state
input_uncoalesced._coalesced_(False)
self._test_asin_arcsin(input_uncoalesced, coalesced)
@coalescedonoff
@dtypes(torch.double)
def test_mv(self, device, dtype, coalesced):
def test_shape(di, dj, dk, nnz):
x, _, _ = self._gen_sparse(2, nnz, [di, dj], dtype, device, coalesced)
t = torch.randn(dk, dtype=dtype, device=device)
res = x.matmul(t)
expected = self.safeToDense(x).matmul(t)
self.assertEqual(res, expected)
test_shape(10, 100, 100, 20)
test_shape(100, 1000, 1000, 20)
test_shape(64, 10000, 10000, 20)
test_shape(0, 100, 100, 0)
test_shape(10, 0, 0, 0)
test_shape(10, 100, 100, 0)
test_shape(10, 100, 100, 20)
with self.assertRaisesRegex(RuntimeError, r"mv: expected self\.size\(-1\) == vec\.size\(-1\)"):
test_shape(10, 100, 10, 20)
with self.assertRaisesRegex(RuntimeError, "mv: two tensor dim should be 2 and 1"):
x, _, _ = self._gen_sparse(2, 20, [10, 100], dtype, device, coalesced)
y, _, _ = self._gen_sparse(2, 20, [10, 100], dtype, device, coalesced)
res = x.mv(y)
@dtypes(*floating_and_complex_types())
def test_sparse_add_coalesce(self, device, dtype):
i = self.index_tensor([[1, 2, 1]], device=device)
v = torch.tensor([3, 4, 5], dtype=dtype, device=device)
x = self.sparse_tensor(i, v, torch.Size([3]))
y = self.sparse_tensor(i, v, torch.Size([3]))
z = x + y
self.assertFalse(z._indices().numel() != 2 and z.is_coalesced())
i = self.index_tensor([[1, 2, 1]], device=device)
v = torch.empty([3, 0], dtype=dtype, device=device)
x = self.sparse_tensor(i, v, torch.Size([3, 0]))
y = self.sparse_tensor(i, v, torch.Size([3, 0]))
z = x + y
self.assertFalse(z._indices().numel() != 2 and z.is_coalesced())
@onlyCUDA
def test_storage_not_null(self, device):
x = torch.sparse_coo_tensor((2,), dtype=torch.float32, device=device)
self.assertNotEqual(x.get_device(), -1)
x = torch.sparse_coo_tensor((2, 0), dtype=torch.float32, device=device)
self.assertNotEqual(x.get_device(), -1)
@onlyCUDA
@deviceCountAtLeast(2)
def test_same_gpu(self, devices):
def check_device(x, device_id):
self.assertEqual(x.get_device(), device_id)
self.assertEqual(x._values().get_device(), device_id)
self.assertEqual(x._indices().get_device(), device_id)
dev1, dev2 = devices[0], devices[1]
i = self.index_tensor([[2]], device=dev2)
v = torch.tensor([5], device=dev2)
x = self.sparse_tensor(i, v, torch.Size([3]), device=1)
check_device(x, 1)
i = self.index_tensor([[2]], device=dev2)
v = torch.empty(1, 0, device=dev2)
x = self.sparse_tensor(i, v, torch.Size([3, 0]), device=1)
check_device(x, 1)
x = self.sparse_empty(3, device=1)
check_device(x, 1)
x = self.sparse_empty(3, 0, device=1)
check_device(x, 1)
def _test_new_device(self, size, device=torch.cuda):
with torch.cuda.device(device):
x = torch.sparse_coo_tensor(size, device='cuda', dtype=torch.float64)
self.assertEqual(x.get_device(), device)
x1 = x.new()
x2 = x.new(2, 3)
self.assertEqual(x1.get_device(), device)
self.assertEqual(x2.get_device(), device)
@onlyCUDA
def test_new_device_single_gpu(self):
self._test_new_device((), 0)
self._test_new_device((30, 20), 0)
self._test_new_device((30, 20, 10), 0)
self._test_new_device((30, 20, 10, 0), 0)
@onlyCUDA
@unittest.skipIf(not TEST_MULTIGPU, "only one GPU detected")
def test_new_device_multi_gpu(self):
self._test_new_device((), 1)
self._test_new_device((30, 20), 1)
self._test_new_device((30, 20, 10), 1)
self._test_new_device((30, 20, 10, 0), 1)
@coalescedonoff
@dtypes(torch.double, torch.cdouble)
def test_new(self, device, dtype, coalesced):
def test_shape(sparse_dims, nnz, with_size):
x, indices, values = self._gen_sparse(sparse_dims, nnz, with_size, dtype, device, coalesced)
if not x.is_cuda:
# CUDA sparse tensors currently requires the size to be
# specified if nDimV > 0
out = x.new(indices, values).coalesce()
x_c = x.coalesce()
self.assertEqual((out.indices(), out.values()), (x_c.indices(), x_c.values()))
self.assertEqual(x.new(indices, values, x.size()), x)
test_shape(3, 10, 100)
test_shape(3, 0, [100, 100, 0])
@onlyCPU # not really, but we only really want to run this once
@dtypes(torch.float64, torch.float32, torch.float16, torch.cfloat, torch.cdouble)
def test_factory(self, device, dtype):
for test_empty_tensor in [True, False]:
if test_empty_tensor:
default_size = torch.Size([1, 3, 0])
size = torch.Size([3, 3, 0])
else:
default_size = torch.Size([1, 3])
size = torch.Size([3, 3])
for include_size in [True, False]:
for use_tensor_idx in [True, False]:
for use_tensor_val in [True, False]:
for use_cuda in ([False] if not torch.cuda.is_available() else [True, False]):
# have to include size with cuda sparse tensors
include_size = include_size or use_cuda
long_dtype = torch.int64
device = torch.device('cpu') if not use_cuda else \
torch.device(torch.cuda.device_count() - 1)
indices = torch.tensor(([0], [2]), dtype=long_dtype) if use_tensor_idx else ([0], [2])
if test_empty_tensor:
values = torch.empty(1, 0).to(dtype)
else:
if use_tensor_val:
values = torch.tensor([1.], dtype=dtype)
else:
values = 1.
if include_size:
sparse_tensor = torch.sparse_coo_tensor(indices, values, size, dtype=dtype,
device=device, requires_grad=True)
else:
sparse_tensor = torch.sparse_coo_tensor(indices, values, dtype=dtype,
device=device, requires_grad=True)
self.assertEqual(indices, sparse_tensor._indices())
self.assertEqual(values, sparse_tensor._values())
self.assertEqual(size if include_size else default_size, sparse_tensor.size())
self.assertEqual(dtype, sparse_tensor.dtype)
if use_cuda:
self.assertEqual(device, sparse_tensor._values().device)
self.assertEqual(True, sparse_tensor.requires_grad)
@dtypes(torch.double, torch.cdouble)
def test_factory_size_check(self, device, dtype):
indices = self.index_tensor([[1, 2],
[0, 2]], device=device)
values = torch.tensor([.5, .5], dtype=dtype, device=device)
sizes = torch.Size([2, 3])
with self.assertRaisesRegex(RuntimeError, "size is inconsistent with indices"):
torch.sparse_coo_tensor(indices, values, sizes, dtype=dtype, device=device)
indices.fill_(-1)
with self.assertRaisesRegex(RuntimeError, "found negative index"):
torch.sparse_coo_tensor(indices, values, sizes, dtype=dtype, device=device)
indices = self.index_tensor([[1, 2],
[0, 2]], device=device)
values = torch.empty([2, 1, 0], dtype=dtype, device=device)
sizes = torch.Size([2, 3, 1, 0])
with self.assertRaisesRegex(RuntimeError, "size is inconsistent with indices"):
torch.sparse_coo_tensor(indices, values, sizes, dtype=dtype, device=device)
indices = self.index_tensor([[1, 2],
[0, 2]], device=device)
values = torch.empty([2, 2, 2], dtype=dtype, device=device)
sizes = torch.Size([0, 0, 2, 2])
with self.assertRaisesRegex(RuntimeError, "size is inconsistent with indices"):
torch.sparse_coo_tensor(indices, values, sizes, dtype=dtype, device=device)
indices = self.index_tensor([[1, 2],
[0, 2]], device=device)
values = torch.tensor([[1, 1, 1], [1, 1, 1]], dtype=dtype, device=device)
sizes = torch.Size([3, 3, 2])
with self.assertRaisesRegex(RuntimeError, "values has incorrect size"):
torch.sparse_coo_tensor(indices, values, sizes, dtype=dtype, device=device)
indices = self.index_tensor([[1, 2],
[0, 2]], device=device)
values = torch.empty([2, 1, 0], dtype=dtype, device=device)
sizes = torch.Size([3, 3, 2, 0])
with self.assertRaisesRegex(RuntimeError, "values has incorrect size"):
torch.sparse_coo_tensor(indices, values, sizes, dtype=dtype, device=device)
def test_factory_empty_indices(self, device):
tensor = torch.sparse_coo_tensor(torch.Size([2, 0]), device=device)
expected_indices = torch.empty((2, 0), dtype=torch.long, device=device)
self.assertEqual(tensor._indices(), expected_indices)
tensor = torch.sparse_coo_tensor(torch.Size([2, 2, 0]), device=device)
expected_indices = torch.empty((3, 0), dtype=torch.long, device=device)
self.assertEqual(tensor._indices(), expected_indices)
tensor = torch.sparse_coo_tensor(torch.Size([2, 2, 0, 0]), device=device)
expected_indices = torch.empty((4, 0), dtype=torch.long, device=device)
self.assertEqual(tensor._indices(), expected_indices)
@dtypes(torch.double, torch.cdouble)
def test_factory_nnz(self, device, dtype):
indices = self.index_tensor([[0]], device=device) # (sparse_dim, nnz): (1, 1)
values = torch.tensor([[1, 1], [1, 1]], dtype=dtype, device=device) # (nnz, ...): (2, 2)
sizes = torch.Size([2, 2])
with self.assertRaisesRegex(RuntimeError, "indices and values must have same nnz"):
torch.sparse_coo_tensor(indices, values, sizes, dtype=dtype, device=device)
indices = self.index_tensor([[0]], device=device) # (sparse_dim, nnz): (1, 1)
values = torch.empty([2, 0], dtype=dtype, device=device) # (nnz, ...): (2, 0)
sizes = torch.Size([2, 0])
with self.assertRaisesRegex(RuntimeError, "indices and values must have same nnz"):
torch.sparse_coo_tensor(indices, values, sizes, dtype=dtype, device=device)
@dtypes(torch.double, torch.cdouble)
def test_factory_nnz_zero(self, device, dtype):
def test_shape(i_shape, v_shape, size, expected_size):
if size:
t = torch.sparse_coo_tensor(torch.empty(i_shape), torch.empty(v_shape), torch.Size(size),
dtype=dtype, device=device)
else:
t = torch.sparse_coo_tensor(torch.empty(i_shape), torch.empty(v_shape), dtype=dtype, device=device)
expected_indices = torch.empty(i_shape, device=device, dtype=torch.int64)
expected_values = torch.empty(v_shape, device=device, dtype=dtype)
expected_size = torch.Size(expected_size)
self.assertEqual(t._indices(), expected_indices)
self.assertEqual(t._values(), expected_values)
self.assertEqual(t.size(), expected_size)
test_shape([1, 0], [0, 2, 4, 0], None, [0, 2, 4, 0])
test_shape([3, 0], [0, 2, 4, 0], None, [0, 0, 0, 2, 4, 0])
test_shape([1, 0], [0, 2, 4, 0], [0, 2, 4, 0], [0, 2, 4, 0])
test_shape([3, 0], [0, 2, 4, 0], [0, 0, 0, 2, 4, 0], [0, 0, 0, 2, 4, 0])
test_shape([3, 0], [0, 2, 4, 0], [1, 2, 3, 2, 4, 0], [1, 2, 3, 2, 4, 0])
@dtypes(torch.double, torch.cdouble)
def test_factory_dense_dim(self, device, dtype):
indices = self.index_tensor([[0]], device=device)
values = torch.tensor([[[1, 1, 1], [1, 1, 1]]], dtype=dtype, device=device)
sizes = torch.Size([1, 3, 4])
with self.assertRaisesRegex(RuntimeError, "values has incorrect size"):
torch.sparse_coo_tensor(indices, values, sizes)
indices = self.index_tensor([[0]], device=device)
values = torch.empty([1, 2, 3, 0], dtype=dtype, device=device)
sizes = torch.Size([1, 3, 4, 0])
with self.assertRaisesRegex(RuntimeError, "values has incorrect size"):
torch.sparse_coo_tensor(indices, values, sizes)
@onlyCPU
@dtypes(torch.float16, torch.float32, torch.float64, torch.cfloat, torch.cdouble, torch.int64)
def test_factory_type_inference(self, device, dtype):
t = torch.sparse_coo_tensor(torch.tensor(([0], [2])), torch.tensor([1.], dtype=dtype))
self.assertEqual(dtype, t.dtype)
t = torch.sparse_coo_tensor(torch.tensor(([0], [2])), torch.tensor([1]))
self.assertEqual(torch.int64, t.dtype)
t = torch.sparse_coo_tensor(torch.tensor(([0], [2])), torch.HalfTensor(1, 0))
self.assertEqual(torch.float16, t.dtype)
t = torch.sparse_coo_tensor(torch.tensor(([0], [2])), torch.FloatTensor(1, 0))
self.assertEqual(torch.float32, t.dtype)
t = torch.sparse_coo_tensor(torch.tensor(([0], [2])), torch.DoubleTensor(1, 0))
self.assertEqual(torch.float64, t.dtype)
t = torch.sparse_coo_tensor(torch.tensor(([0], [2])), torch.LongTensor(1, 0))
self.assertEqual(torch.int64, t.dtype)
@onlyCUDA
def test_factory_device_type_inference(self, device):
# both indices/values are CUDA
cpu_cuda = ('cpu', 'cuda')
cpu_cuda_none = cpu_cuda + (None,)
for indices_device, values_device, device in itertools.product(cpu_cuda,
cpu_cuda,
cpu_cuda_none):
indices = torch.tensor(([0], [2]), device=indices_device)
values = torch.tensor([1.], device=values_device)
empty_values = torch.empty(1, 0).to(values_device)
shape = (1, 3)
empty_shape = (1, 3, 0)
if device is None and indices_device != values_device:
with self.assertRaises(RuntimeError):
torch.sparse_coo_tensor(indices, values, shape, device=device)
with self.assertRaises(RuntimeError):
torch.sparse_coo_tensor(indices, empty_values, empty_shape, device=device)
else:
t = torch.sparse_coo_tensor(indices, values, shape, device=device)
t_empty = torch.sparse_coo_tensor(indices, empty_values, empty_shape, device=device)
should_be_cuda = (device == 'cuda' or (device is None and values_device == 'cuda'))
self.assertEqual(should_be_cuda, t.is_cuda)
self.assertEqual(t.is_cuda, t_empty.is_cuda)
@onlyCPU
def test_factory_copy(self, device):
def test_tensor(indices, values, indices_equal, values_equal):
sparse_tensor = torch.sparse_coo_tensor(indices, values, dtype=torch.float64, device=device)
if indices_equal:
self.assertEqual(indices.data_ptr(), sparse_tensor._indices().data_ptr())
else:
self.assertNotEqual(indices.data_ptr(), sparse_tensor._indices().data_ptr())
if values_equal:
self.assertEqual(values.data_ptr(), sparse_tensor._values().data_ptr())
else:
self.assertNotEqual(values.data_ptr(), sparse_tensor._values().data_ptr())
# both correct
indices = torch.tensor(([0], [2]), dtype=torch.int64)
values = torch.tensor([1.], dtype=torch.float64)
test_tensor(indices, values, True, True)
indices = torch.tensor(([0], [2]), dtype=torch.int64)
values = torch.DoubleTensor(1, 0)
test_tensor(indices, values, True, True)
# only indices correct
indices = torch.tensor(([0], [2]), dtype=torch.int64)
values = torch.tensor([1.], dtype=torch.float32)
test_tensor(indices, values, True, False)
indices = torch.tensor(([0], [2]), dtype=torch.int64)
values = torch.tensor([1.], dtype=torch.float16)
test_tensor(indices, values, True, False)
indices = torch.tensor(([0], [2]), dtype=torch.int64)
values = torch.FloatTensor(1, 0)
test_tensor(indices, values, True, True) # An empty tensor's data_ptr is always equal to 0
# only values correct
indices = torch.tensor(([0], [2]), dtype=torch.int32)
values = torch.tensor([1.], dtype=torch.float64)
test_tensor(indices, values, False, True)
indices = torch.tensor(([0], [2]), dtype=torch.int32)
values = torch.DoubleTensor(1, 0)
test_tensor(indices, values, False, True)
# neither correct
indices = torch.tensor(([0], [2]), dtype=torch.int32)
values = torch.tensor([1.], dtype=torch.float32)
test_tensor(indices, values, False, False)
indices = torch.tensor(([0], [2]), dtype=torch.int32)
values = torch.FloatTensor(1, 0)
test_tensor(indices, values, False, True) # An empty tensor's data_ptr is always equal to 0
# complex support
indices = torch.tensor(([0], [2]), dtype=torch.int64)
values = make_tensor([1, ], dtype=torch.cdouble, device=device)
test_tensor(indices, values, True, False)
indices = torch.tensor(([0], [2]), dtype=torch.int32)
values = make_tensor([1, 1], dtype=torch.cdouble, device=device)
test_tensor(indices, values, False, False)
@onlyCPU # just run once, we test both cpu and cuda
def test_legacy_new_device(self, device):
i = torch.tensor([[0, 1, 1], [2, 0, 2]])
v = torch.tensor([3., 4., 5.])
size = torch.Size([2, 3])
x = torch.sparse_coo_tensor(i, v, size, device='cpu')
self.assertRaises(RuntimeError, lambda: x.new(device='cuda'))
self.assertRaises(RuntimeError, lambda: x.new(i, v, device='cuda'))
self.assertRaises(RuntimeError, lambda: x.new(i, v, size, device='cuda'))
self.assertRaises(RuntimeError, lambda: x.new(torch.Size([2, 3, 4]), device='cuda'))
if torch.cuda.is_available():
x = torch.sparse_coo_tensor(i, v, size, device='cuda')
self.assertRaises(RuntimeError, lambda: x.new(device='cpu'))
self.assertRaises(RuntimeError, lambda: x.new(i, v, device='cpu'))
self.assertRaises(RuntimeError, lambda: x.new(i, v, size, device='cpu'))
self.assertRaises(RuntimeError, lambda: x.new(torch.Size([2, 3, 4]), device='cpu'))
def test_legacy_new(self, device):
i = torch.tensor([[0, 1, 1], [2, 0, 2]])
v = torch.tensor([3., 4., 5.])
size = torch.Size([2, 3])
s = torch.sparse_coo_tensor(i, v, size)
self.assertEqual(torch.sparse_coo, s.new(device='cpu').layout)
self.assertRaises(TypeError, lambda: s.new(v.untyped_storage()))
self.assertRaises(TypeError, lambda: s.new(v))
self.assertEqual(torch.sparse_coo, s.new(torch.Size([2, 3])).layout)
self.assertRaises(TypeError, lambda: s.new([6]))
@onlyCPU # not really, but we only really want to run this once
def test_dtypes(self, device):
all_sparse_dtypes = all_types_and_complex_and(torch.half, torch.bool, torch.bfloat16)
do_test_dtypes(self, all_sparse_dtypes, torch.sparse_coo, torch.device('cpu'))
if torch.cuda.is_available():
do_test_dtypes(self, all_sparse_dtypes, torch.sparse_coo, torch.device('cuda:0'))
def _test_empty_full(self, device, dtype, requires_grad):
shape = (2, 3)
layout = torch.sparse_coo
def check_value(tensor, value=None, dtype=dtype, requires_grad=requires_grad):
self.assertEqual(shape, tensor.shape)
self.assertIs(dtype, tensor.dtype)
self.assertIs(layout, tensor.layout)
self.assertEqual(tensor.requires_grad, requires_grad)
if tensor.is_cuda and device is not None:
self.assertEqual(device, tensor.device)
if value is not None:
fill = tensor.empty(shape, dtype=dtype).fill_(value)
self.assertEqual(tensor, fill)
v = torch.sparse_coo_tensor(shape, dtype=dtype, device=device, requires_grad=requires_grad)
check_value(v)
out = v.new()
check_value(torch.zeros(shape, out=out, device=device, requires_grad=requires_grad))
int64_dtype = torch.int64
check_value(v.new_empty(shape), requires_grad=False)
check_value(v.new_empty(shape, dtype=int64_dtype, device=device, requires_grad=False),
dtype=int64_dtype, requires_grad=False)
check_value(torch.empty_like(v), requires_grad=False)
check_value(torch.empty_like(v, dtype=int64_dtype, layout=layout, device=device, requires_grad=False),
dtype=int64_dtype, requires_grad=False)
@onlyCPU # not really, but we only really want to run this once
@dtypes(*all_types_and_complex_and(torch.half, torch.bool, torch.bfloat16))
@parametrize('requires_grad', (True, False))
def test_empty_full(self, device, dtype, requires_grad):
if requires_grad and not (dtype.is_floating_point or dtype.is_complex):
self.skipTest(f'requires_grad==True requires float or complex dtype, got {dtype}')
self._test_empty_full(device, dtype, requires_grad)
if torch.cuda.is_available():
self._test_empty_full(None, dtype, requires_grad)
self._test_empty_full(torch.device('cuda:0'), dtype, requires_grad)
def test_is_sparse(self, device):
x = torch.randn(3, 3)
self.assertFalse(x.is_sparse)
x = torch.randn(3, 3, 0)
self.assertFalse(x.is_sparse)
x = self.sparse_empty(1, 0, device=device)
self.assertTrue(x.is_sparse)
def test_resize_as(self, device):
def do_test(t):
y = t.new().resize_as_(t).zero_()
self.assertEqual(y.shape, t.shape)
# Check that y can be added to t. Currently, this requires that
# sparse_dim and dense_dim match.
self.assertEqual(t, t + y)
do_test(self.sparse_empty([3, 0], device=device))
do_test(self.sparse_empty([3, 3], device=device))
def _test_resize_shape(self, x_i, x_v, x_size, y_i, y_v, y_size, dtype, device):
x_v_numel = torch.zeros(x_v).numel()
x = torch.sparse_coo_tensor(torch.zeros(x_i),
torch.arange(x_v_numel).resize_(x_v).to(torch.float),
torch.Size(x_size), dtype=dtype, device=device)
x_dense = x.to_dense()
y = torch.sparse_coo_tensor(torch.zeros(y_i),
torch.ones(y_v).to(torch.float),
torch.Size(y_size), dtype=dtype, device=device)
y_dense = y.to_dense()
x.resize_as_(y)
x_dense.resize_as_(y_dense)
self.assertEqual(x.shape, y.shape)
self.assertEqual(x.sparse_dim(), y.sparse_dim())
self.assertEqual(x.dense_dim(), y.dense_dim())
self.assertEqual(x.shape, x_dense.shape)
self.assertEqual(y.shape, y_dense.shape)
# Here we make sure that the original data are preserved after resizing
self.assertEqual(x.to_dense().view(-1)[0:x_v_numel].view(x_v),
x_dense.view(-1)[0:x_v_numel].view(x_v))
@dtypes(torch.double, torch.cdouble)
def test_resize(self, device, dtype):
# 1. Expand the size of some dense dimensions [Supported]
self._test_resize_shape([1, 1], [1, 2, 3], [2, 2, 3],
[1, 1], [1, 2, 4], [2, 2, 4],
dtype=dtype, device=device)
self._test_resize_shape([1, 1], [1, 2, 0], [2, 2, 0],
[1, 1], [1, 2, 4], [2, 2, 4],
dtype=dtype, device=device)
# 2. Expand the size of some sparse dimensions [Supported]
self._test_resize_shape([1, 1], [1, 2, 3], [2, 2, 3],
[1, 1], [1, 2, 3], [4, 2, 3],
dtype=dtype, device=device)
# 3. Change the shapes of both sparse and dense dimensions when nnz is zero [Supported]
self._test_resize_shape([1, 0], [0, 2, 3], [2, 2, 3],
[2, 0], [0, 2, 4, 5], [1, 1, 2, 4, 5],
dtype=dtype, device=device)
self._test_resize_shape([1, 0], [0, 2, 3], [2, 2, 3],
[2, 0], [0, 2, 4, 0], [1, 1, 2, 4, 0],
dtype=dtype, device=device)
# 4. Add dims to dense dimensions [Not Supported]
with self.assertRaisesRegex(RuntimeError, "changing the number of dense dimensions"):
self._test_resize_shape([1, 1], [1, 2, 3], [2, 2, 3],
[1, 1], [1, 2, 3, 4], [2, 2, 3, 4],
dtype=dtype, device=device)
with self.assertRaisesRegex(RuntimeError, "changing the number of dense dimensions"):
self._test_resize_shape([1, 1], [1, 2, 3], [2, 2, 3],
[1, 1], [1, 2, 3, 0], [2, 2, 3, 0],
dtype=dtype, device=device)
# 5. Remove dims from dense dimensions [Not Supported]
with self.assertRaisesRegex(RuntimeError, "changing the number of dense dimensions"):
self._test_resize_shape([1, 1], [1, 2, 3], [2, 2, 3],
[1, 1], [1, 2], [2, 2],
dtype=dtype, device=device)
# 6. Change the number of sparse dimensions on a non-empty sparse tensor [Not Supported]
with self.assertRaisesRegex(RuntimeError, "changing the number of sparse dimensions"):
self._test_resize_shape([1, 1], [1, 2, 3], [2, 2, 3],
[2, 1], [1, 2, 3], [1, 2, 2, 3],
dtype=dtype, device=device)
# 7. Shrink the size of some sparse dimensions on a non-empty sparse tensor [Not Supported]
with self.assertRaisesRegex(RuntimeError, "shrinking the size of sparse dimensions"):
self._test_resize_shape([1, 1], [1, 2, 3], [2, 2, 3],
[1, 1], [1, 2, 3], [1, 2, 3],
dtype=dtype, device=device)
# 8. Shrink the size of some dense dimensions on a non-empty sparse tensor [Not Supported]
with self.assertRaisesRegex(RuntimeError, "shrinking the size of dense dimensions"):
self._test_resize_shape([1, 1], [1, 2, 3], [2, 2, 3],
[1, 1], [1, 2, 2], [2, 2, 2],
dtype=dtype, device=device)
with self.assertRaisesRegex(RuntimeError, "shrinking the size of dense dimensions"):
self._test_resize_shape([1, 1], [1, 2, 3], [2, 2, 3],
[1, 1], [1, 2, 0], [2, 2, 0],
dtype=dtype, device=device)
def test_is_nonzero(self, device):
self.assertTrue(torch.sparse_coo_tensor(([0],), 1., (1,), device=device).is_nonzero())
self.assertFalse(torch.sparse_coo_tensor(([0],), 0., (1,), device=device).is_nonzero())
self.assertFalse(torch.sparse_coo_tensor(([0], [0]), 0., (1, 1), device=device).is_nonzero())
self.assertFalse(torch.sparse_coo_tensor(([0, 0],), (0., 0.), (1,), device=device).is_nonzero())
self.assertFalse(torch.sparse_coo_tensor(([0, 0],), (-1., 1.), (1,), device=device).is_nonzero())
# scalar sparse tensor
self.assertTrue(torch.sparse_coo_tensor(torch.zeros(0, 1), 12.3, [], device=device).is_nonzero())
with self.assertRaisesRegex(RuntimeError, "Boolean value of Tensor with no values is ambiguous"):
torch.sparse_coo_tensor(([0, 1],), torch.empty(2, 0), (4, 0), device=device).is_nonzero()
self.assertTrue(torch.sparse_coo_tensor(([0],), 2.3 - 4.5j, (1,), dtype=torch.cfloat, device=device)
.is_nonzero())
self.assertTrue(torch.sparse_coo_tensor(([0],), 2.3 - 4.5j, (1,), dtype=torch.cdouble, device=device)
.is_nonzero())
self.assertFalse(torch.sparse_coo_tensor(([0],), 0. + 0j, (1,), dtype=torch.cfloat, device=device)
.is_nonzero())
self.assertFalse(torch.sparse_coo_tensor(([0],), 0. + 0j, (1,), dtype=torch.cdouble, device=device)
.is_nonzero())
@dtypes(torch.double, torch.cdouble)
def test_change_tensor_metadata(self, device, dtype):
i = self.index_tensor([[0], [1]], device=device)
v = torch.tensor([[3, 4, 5]], dtype=dtype, device=device)
t = torch.sparse_coo_tensor(i, v, torch.Size([1, 2, 3]), dtype=dtype, device=device)
i.resize_(2, 3)
v.resize_(4, 5)
self.assertEqual(list(t.coalesce().indices().size()), [2, 1])
self.assertEqual(list(t.coalesce().values().size()), [1, 3])
i = self.index_tensor([[0], [1]], device=device)
v = torch.tensor([[3, 4, 5]], dtype=dtype, device=device)
t = torch.sparse_coo_tensor(i, v, torch.Size([1, 2, 3]))
i.resize_as_(self.index_tensor([0, 1], device=device))
v.resize_as_(torch.tensor([3, 4, 5], dtype=dtype, device=device))
self.assertEqual(list(t.coalesce().indices().size()), [2, 1])
self.assertEqual(list(t.coalesce().values().size()), [1, 3])
i = self.index_tensor([[0], [1]], device=device)
v = torch.tensor([[3, 4, 5]], dtype=dtype, device=device)
t = torch.sparse_coo_tensor(i, v, torch.Size([1, 2, 3]))
i.as_strided_((2, 1), (1, 1))
v.as_strided_((1, 3), (1, 1))
self.assertEqual(list(t.coalesce().indices().size()), [2, 1])
self.assertEqual(list(t.coalesce().values().size()), [1, 3])
i = self.index_tensor([[0], [1]], device=device)
v = torch.tensor([[3, 4, 5]], dtype=dtype, device=device)
t = torch.sparse_coo_tensor(i, v, torch.Size([1, 2, 3]))
i.set_(self.index_tensor([0, 1], device=device))
v.set_(torch.tensor([3, 4, 5], dtype=dtype, device=device))
self.assertEqual(list(t.coalesce().indices().size()), [2, 1])
self.assertEqual(list(t.coalesce().values().size()), [1, 3])
i = self.index_tensor([[0], [1]], device=device)
v = torch.tensor([[3, 4, 5]], dtype=dtype, device=device)
t = torch.sparse_coo_tensor(i, v, torch.Size([1, 2, 3]))
i.transpose_(0, 1)
v.transpose_(0, 1)
self.assertEqual(list(t.coalesce().indices().size()), [2, 1])
self.assertEqual(list(t.coalesce().values().size()), [1, 3])
@coalescedonoff
@dtypes(torch.double)
def test_pickle(self, device, dtype, coalesced):
import pickle
shape_sparse_dim_nnz = [
((), 0, 2),
((0,), 0, 10),
((2,), 0, 3),
((100, 3), 1, 3),
((100, 20, 3), 2, 0),
((10, 0, 3), 0, 3),
((10, 0, 3), 0, 0),
]
for shape, sparse_dim, nnz in shape_sparse_dim_nnz:
indices_shape = torch.Size((sparse_dim, nnz))
values_shape = torch.Size((nnz,) + shape[sparse_dim:])
indices = torch.arange(indices_shape.numel(), dtype=self.index_tensor(0).dtype,
device=device).view(indices_shape)
for d in range(sparse_dim):
indices[d].clamp_(max=(shape[d] - 1)) # make it valid index
if not coalesced and indices.numel() > 0:
indices[:, -1] = indices[:, 0] # make it uncoalesced
values_numel = values_shape.numel()
values = torch.arange(values_numel, dtype=dtype,
device=device).view(values_shape).div_(values_numel / 2.)
sp_tensor = self.sparse_tensor(indices, values, shape)
serialized = pickle.dumps(sp_tensor)
sp_tensor_loaded = pickle.loads(serialized)
self.assertEqual(sp_tensor, sp_tensor_loaded)
def test_any(self, device):
t = torch.sparse_coo_tensor(torch.tensor(([0, 0], [2, 0])), torch.tensor([False, False]), device=device)
t_any = torch.tensor(False)
self.assertEqual(torch.any(t), t_any)
t = torch.sparse_coo_tensor(torch.tensor(([0, 0], [2, 0])), torch.tensor([True, False]), device=device)
t_any = torch.tensor(True)
self.assertEqual(torch.any(t), t_any)
def test_isnan(self, device):
t = torch.sparse_coo_tensor(torch.tensor(([0, 0], [0, 2])), torch.tensor([1, 4]), device=device)
t_nan = torch.sparse_coo_tensor(torch.tensor(([0, 0], [0, 2])), torch.tensor([False, False]), device=device)
self.assertEqual(torch.isnan(t).int(), t_nan.int())
t = torch.sparse_coo_tensor(torch.tensor(([0, 0], [0, 2])), torch.tensor([1, float("nan")]), device=device)
t_nan = torch.sparse_coo_tensor(torch.tensor(([0, 0], [0, 2])), torch.tensor([False, True]), device=device)
self.assertEqual(torch.isnan(t).int(), t_nan.int())
@coalescedonoff
@dtypes(torch.float32, torch.float64)
def test_div_rounding_mode(self, device, dtype, coalesced):
sparse, _, _ = self._gen_sparse(2, 10, (10, 10), dtype,
device, coalesced)
dense = self.safeToDense(sparse)
for mode in (None, 'floor', 'trunc'):
actual = sparse.div(-2, rounding_mode=mode)
expect = dense.div(-2, rounding_mode=mode)
self.assertEqual(self.safeToDense(actual), expect)
# Test inplace
actual = sparse.clone().div_(-2, rounding_mode=mode)
self.assertEqual(self.safeToDense(actual), expect)
# Test out argument
actual.zero_()
torch.div(sparse, -2, rounding_mode=mode, out=actual)
self.assertEqual(self.safeToDense(actual), expect)
def test_div_by_sparse_error(self, device):
self.assertRaisesRegex(RuntimeError, 'Sparse division requires',
lambda: torch.tensor(1., device=device).to_sparse()
/ torch.tensor(1., device=device).to_sparse())
def test_floor_divide_by_sparse_error(self, device):
self.assertRaisesRegex(RuntimeError, 'Sparse floor division requires',
lambda: torch.tensor(1., device=device).to_sparse()
// torch.tensor(1., device=device).to_sparse())
@unittest.skipIf(not TEST_NUMPY, "Numpy not found")
@onlyCPU
def test_sparse_to_numpy(self, device):
t = torch.sparse_coo_tensor(torch.tensor(([0, 0], [2, 0])), torch.tensor([1, 4]))
self.assertRaises(TypeError, lambda: t.numpy())
@coalescedonoff
@dtypes(torch.double)
def test_softmax(self, device, dtype, coalesced):
import torch.nn.functional as F
def to_dense(sparse, fill_value=None):
"""
Return dense tensor from a sparse tensor using given fill value.
"""
if fill_value is None or fill_value == 0:
return sparse.to_dense()
sparse = sparse.coalesce()
dense = torch.full(sparse.shape, fill_value, dtype=sparse.dtype, device=sparse.device)
for idx, value in zip(sparse._indices().t(), sparse._values()):
dense[tuple(idx)] = value
return dense
def softmax_to_dense(sparse, dim):
"""Dense softmax of a sparse tensor. Useful only for testing softmax
correctness.
When computing softmax of a sparse tensor, the value of
unspecified items is negative infinity rather than zero so
that
softmax(sparse.to_dense(fill_value=-inf), dim) == softmax(sparse, dim).to_dense()
holds for non-empty lines. One empty lines, the softmax
values are defined as 0 in order to preserve the sparsity
of result.
Note that in PyTorch, ``to_dense`` method does not
implement the ``fill_value`` keyword argument.
"""
dtype = sparse.dtype
device = sparse.device
dense = to_dense(sparse, fill_value=-float('inf'))
r = F.softmax(dense, dim)
# softmax on empty lines results nan, replace with zeros to match the definition
r[r != r] = 0
return r
def sparse_softmax(sparse, dim):
"""Pure Python softmax of a sparse tensor. Assuming -inf for
unspecified sparse tensor data. This is a prototype of
sparse softmax algorithm in Python.
"""
dtype = sparse.dtype
device = sparse.device
# softmax is non-linear operation, so sparse tensors must
# be coalesced.
sparse = sparse.coalesce()
inf = float('inf')
indices = sparse._indices()
values = sparse._values()
if dim < sparse.sparse_dim():
nnz = sparse._nnz()
# compute pool indices
size = sparse.size()
strides = torch.ones((sparse.sparse_dim(), 1), dtype=indices.dtype, device=indices.device)
for i in reversed(range(sparse.sparse_dim() - 1)):
strides[i, 0] = strides[i + 1, 0] * size[i + 1]
strides[dim, 0] = 0
pool = (indices * strides).sum(dim=0)
i2p = {}
for i in range(nnz):
c = int(pool[i])
if c not in i2p:
i2p[c] = len(i2p)
pool[i] = i2p[c]
# compute max
dense_size = tuple(size[sparse.sparse_dim():])
mx = torch.empty((pool.max() + 1,) + dense_size, dtype=dtype, device=device)
mx[:] = -inf
for n in range(nnz):
p = pool[n]
mx[p] = torch.max(mx[p], values[n])
# apply exp to (v - mx) and sum the results
exp_values = torch.empty_like(values)
exp_sums = torch.zeros_like(mx)
for n in range(nnz):
p = pool[n]
v = exp_values[n] = (values[n] - mx[p]).exp()
exp_sums[p] = exp_sums[p] + v
# normalize with the sum of exponents
for n in range(nnz):
p = pool[n]
exp_values[n] = exp_values[n] / exp_sums[p]
return torch.sparse_coo_tensor(indices,
exp_values,
sparse.size(),
dtype=dtype, device=device)
elif dim < sparse.sparse_dim() + sparse.dense_dim():
return torch.sparse_coo_tensor(indices,
F.softmax(values, dim - sparse.sparse_dim() + 1),
sparse.size(),
dtype=dtype, device=device)
else:
raise ValueError(
f'`dim(={dim})` must be smaller than `sparse_dim(={sparse.sparse_dim()}) + dense_dim(={sparse.dense_dim()})`')
def softmax_jacobian_analytic(x, dim):
"""Return Jacobian of softmax using analytic formula
D_jS_i = S_i * (1[i==j] - S_j).
where S = softmax(x, dim), x is dense tensor, i,j in
range(x.shape[dim]).
"""
y = F.softmax(x, dim)
y[y != y] = 0 # replace nan-s with zeros
J = torch.zeros((x.shape[dim],) + tuple(x.shape), dtype=x.dtype, device=x.device)
si = [slice(None)] * len(y.shape)
sj = [slice(None)] * len(y.shape)
s = [slice(None)] * len(J.shape)
for i in range(y.shape[dim]):
si[dim] = i
s[dim + 1] = i
yi = y[tuple(si)]
for j in range(y.shape[dim]):
sj[dim] = j
s[0] = j
if i == j:
J[tuple(s)] = yi * (1 - yi)
else:
yj = y[tuple(sj)]
J[tuple(s)] = - yi * yj
sj[dim] = slice(None)
si[dim] = slice(None)
s[dim + 1] = slice(None)
return J
def softmax_jacobian_autograd(x, dim, log=False):
"""Return Jacobian of softmax using PyTorch autograd feature.
x can be dense or sparse tensor.
"""
import itertools
if x.is_sparse:
x = x.coalesce()
dtype = x.dtype
device = x.device
shape = tuple(x.shape)
J = torch.zeros((shape[dim],) + shape, dtype=dtype, device=device)
for i in range(shape[dim]):
if x.is_sparse:
sparse_dim = x.sparse_dim()
dense_dim = x.dense_dim()
if dim < sparse_dim:
ranges = []
for j, sz in enumerate(shape[:sparse_dim]):
if dim == j:
ranges.append([i])
else:
ranges.append(list(range(sz)))
indices = torch.tensor(list(itertools.product(*ranges)), dtype=torch.long, device=device).t()
values = torch.ones((indices.shape[1],) + shape[sparse_dim:], dtype=dtype, device=device)
else:
ranges = []
for j, sz in enumerate(shape[:sparse_dim]):
ranges.append(list(range(sz)))
indices = torch.tensor(list(itertools.product(*ranges)), dtype=torch.long, device=device).t()
values = torch.zeros((indices.shape[1],) + shape[sparse_dim:], dtype=dtype, device=device)
sv = [slice(None)] * (dense_dim + 1)
sv[dim - sparse_dim + 1] = i
values[tuple(sv)] = 1
v = torch.sparse_coo_tensor(indices, values, shape, dtype=dtype, device=device)
else:
v = torch.zeros_like(x)
sv = [slice(None)] * len(v.shape)
sv[dim] = i
v[tuple(sv)] = 1
x_ = x.clone()
x_.requires_grad_(True)
if log:
if x_.is_sparse:
y = torch.sparse.log_softmax(x_, dim)
else:
y = F.log_softmax(x_, dim)
else:
if x_.is_sparse:
y = torch.sparse.softmax(x_, dim)
else:
y = F.softmax(x_, dim)
# replace nan-s with zeros
y.data[y != y] = 0
y.backward(v)
g = x_.grad
if not g.is_sparse:
# replace nan-s with zeros
g.data[g != g] = 0
J[i] = g.to_dense() if g.is_sparse else g
return J
@skipIfTorchDynamo("https://github.com/pytorch/torchdynamo/issues/1166")
def test_op(sparse_dims, nnz, with_size, coalesced):
if isinstance(with_size, Number):
with_size = [with_size] * sparse_dims
x, i, v = self._gen_sparse(sparse_dims, nnz, with_size, dtype, device, coalesced)
def sparse_log(x):
return torch.sparse_coo_tensor(x._indices(), x._values().log(),
x.size(), dtype=x.dtype, device=x.device)
# Check dim out of bounds
with self.assertRaisesRegex(IndexError, r"Dimension out of range"):
torch.sparse.softmax(x, x.dim())
with self.assertRaisesRegex(IndexError, r"Dimension out of range"):
torch.sparse.softmax(x, -x.dim() - 1)
for dim in range(x.dim()):
# Check sparse softmax definition
# check Python sparse softmax
y = sparse_softmax(x, dim)
r1 = softmax_to_dense(x, dim)
r2 = y.to_dense()
self.assertEqual(r1, r2)
# check C++ sparse softmax
for d in (dim, dim - x.dim()):
y1 = torch.sparse.softmax(x, d)
self.assertEqual(y, y1)
# check C++ sparse log_softmax
ly1 = torch.sparse.log_softmax(x, d)
self.assertEqual(ly1, sparse_log(y1))
# Check autograd support on sparse softmax
# check softmax Jacobian definition for dense input
x1 = to_dense(x, fill_value=float('-inf'))
J = softmax_jacobian_analytic(x1, dim)
assert J.shape[0] == x.shape[dim]
assert J.shape[dim + 1] == x.shape[dim]
# check softmax Jacobian from autograd, dense input
J2 = softmax_jacobian_autograd(x1, dim)
self.assertEqual(J, J2)
# check softmax Jacobian from autograd, sparse input
J3 = softmax_jacobian_autograd(x, dim)
self.assertEqual(J, J3)
'''
y = softmax(x, dim)
z = log(y) = log_softmax(x, dim)
Dy/Dx = J
Dz/Dx = Dz/Dy Dy/Dx = 1/y * J
=> J = J_log * y
'''
# log_softmax Jacobian from autograd, dense input
J2_log = softmax_jacobian_autograd(x1, dim, log=True)
# log_softmax Jacobian from autograd, sparse input
J3_log = softmax_jacobian_autograd(x, dim, log=True)
J = J.transpose(0, dim + 1)
J2_log = J2_log.transpose(0, dim + 1)
J3_log = J3_log.transpose(0, dim + 1)
self.assertEqual(J, J2_log * r1)
self.assertEqual(J, J3_log * r1)
if dim == 0:
# check dtype argument
other_dtype = torch.float32
y2 = torch.sparse.softmax(x, dim, dtype=other_dtype)
self.assertEqual(y2.dtype, other_dtype)
self.assertEqual(y2, y1.type(other_dtype))
ly2 = torch.sparse.log_softmax(x, dim, dtype=other_dtype)
self.assertEqual(ly2.dtype, other_dtype)
self.assertEqual(ly2, ly1.type(other_dtype))
test_op(1, 10, [3], coalesced)
test_op(1, 10, [2, 3], coalesced)
test_op(1, 10, [3, 2], coalesced)
test_op(2, 10, [2, 3, 4], coalesced)
test_op(2, 10, [3, 4], coalesced)
test_op(2, 5, [5, 4], coalesced)
test_op(2, 10, [3, 4, 2], coalesced)
test_op(3, 10, [3, 4, 2], coalesced)
test_op(3, 100, [3, 4, 2], coalesced)
test_op(3, 100, [3, 4, 2, 3], coalesced)
test_op(3, 100, [3, 4, 2, 3, 5, 2], coalesced)
test_op(4, 100, [3, 4, 2, 3, 5, 2], coalesced)
def _check_zero_nnz_softmax_op(self, func, ndim, device, dtype):
# create a sparse tensor with shape (0,..., 3) it has no materialize values
t = torch.sparse_coo_tensor([[] for _ in range(ndim)], [], (0,) * (ndim - 1) + (3,), device=device, dtype=dtype)
out = func(t, 0)
self.assertEqual(out, torch.zeros_like(t))
# gradient
t = t.requires_grad_()
gradcheck(lambda x: func(x, 0).to_dense(), (t,), masked=True)
@dtypes(torch.double, torch.float)
@unittest.skipIf(TEST_WITH_CROSSREF, "generator unsupport triggers assertion error")
def test_softmax_zero_nnz(self, device, dtype):
self._check_zero_nnz_softmax_op(torch.sparse.softmax, 1, device, dtype)
self._check_zero_nnz_softmax_op(torch.sparse.softmax, 10, device, dtype)
@dtypes(torch.double, torch.float)
@unittest.skipIf(TEST_WITH_CROSSREF, "generator unsupport triggers assertion error")
def test_log_softmax_zero_nnz(self, device, dtype):
self._check_zero_nnz_softmax_op(torch.sparse.log_softmax, 1, device, dtype)
self._check_zero_nnz_softmax_op(torch.sparse.log_softmax, 10, device, dtype)
# TODO: Check after why ROCm's cusparseXcsrgemm2Nnz function doesn't return the same nnz value as CUDA
@skipIfRocm
@coalescedonoff
@dtypes(*floating_and_complex_types())
@dtypesIfCUDA(*floating_types_and(*[torch.half] if SM53OrLater else [],
*[torch.bfloat16] if SM80OrLater else [],
torch.complex64,
*[torch.complex128] if CUSPARSE_SPMM_COMPLEX128_SUPPORTED else []))
@unittest.skipIf(TEST_WITH_CROSSREF, "not working with fake tensor")
@precisionOverride({torch.bfloat16: 1e-2, torch.float16: 1e-2, torch.complex64: 1e-2, torch.float32: 1e-2})
def test_sparse_matmul(self, device, dtype, coalesced):
"""
This function test `torch.sparse.mm` when both the mat1 and mat2 are sparse tensors.
"""
def ref_sparse_mm(a, b):
return a.to_dense() @ b.to_dense()
def grad_with_custom_sparsity_pattern_test_helper(sparse_dims, nnz, shape_a, shape_b):
def test_grad_dense(a_s, b_s, g_s):
a = a_s.to_dense().detach()
b = b_s.to_dense().detach()
g = g_s.to_dense().detach()
a.requires_grad_(True)
b.requires_grad_(True)
c = a @ b
c.backward(g)
return a.grad.sparse_mask(a_s.coalesce()), b.grad.sparse_mask(b_s.coalesce())
a, _, _ = self._gen_sparse(sparse_dims, nnz, shape_a, dtype, device, coalesced)
b, _, _ = self._gen_sparse(sparse_dims, nnz, shape_b, dtype, device, coalesced)
a.requires_grad_(True)
b.requires_grad_(True)
c = torch.sparse.mm(a, b)
c2 = c.to_dense().detach()
c2 = torch.rand_like(c2)
g = c2.sparse_mask(c.coalesce())
c.backward(g)
a_grad, b_grad = test_grad_dense(a, b, g)
# We convert grad to dense since dense and sparse mm
# implementations handle materialized zeroes differently.
self.assertEqual(a.grad.to_dense(), a_grad.to_dense())
self.assertEqual(b.grad.to_dense(), b_grad.to_dense())
def test_sparse_matmul(sparse_dims, nnz, shape_a, shape_b):
a, i_a, v_a = self._gen_sparse(sparse_dims, nnz, shape_a, dtype, device, coalesced)
b, i_b, v_b = self._gen_sparse(sparse_dims, nnz, shape_b, dtype, device, coalesced)
# dense implementation
r1 = ref_sparse_mm(a, b)
# cpp implementation
r2 = torch.sparse.mm(a, b)
self.assertEqual(r1, r2.to_dense())
# Check result is truly coalesced
self.assertTrue(r2.is_coalesced() and is_coalesced_indices(r2))
if dtype in [torch.double, torch.cdouble]:
a.requires_grad_(True)
b.requires_grad_(True)
# check autograd support on sparse matmul
def fn(D1, D2):
return torch.sparse.mm(D1, D2).to_dense()
if a.is_cuda:
# For cuda, `nondet_tol` is set with `1e-5`
# This is because cuSparse sometimes returns approximate zero values like `~e-323`
# TODO: Check this cuSparse issue.
# This happens when you do chain multiplication `torch.sparse.mm` operations
gradcheck(fn, (a, b), nondet_tol=1e-5, masked=True)
else:
gradcheck(fn, (a, b), masked=True)
grad_with_custom_sparsity_pattern_test_helper(sparse_dims, nnz, shape_a, shape_b)
def test_error_cases():
def fn(sparse_dims, nnz, shape_a, shape_b):
a, i_a, v_a = self._gen_sparse(sparse_dims, nnz, shape_a, dtype, device, coalesced)
b, i_b, v_b = self._gen_sparse(sparse_dims, nnz, shape_b, dtype, device, coalesced)
r2 = torch.sparse.mm(a, b)
# This is not a matrix
self.assertRaises(RuntimeError, lambda: fn(3, 4, [2, 2, 2], [2, 2, 2]))
# Shapes does not
self.assertRaisesRegex(RuntimeError,
r"mat1 and mat2 shapes cannot be multiplied \(2x3 and 4x2\)",
lambda: fn(2, 10, [2, 3], [4, 2]))
def different_dtypes():
a, i_a, v_a = self._gen_sparse(2, 10, [2, 2], dtype, device, coalesced)
b, i_b, v_b = self._gen_sparse(2, 10, [2, 2], dtype, device, coalesced)
r2 = torch.sparse.mm(a.to(torch.float64), a.to(torch.float32))
self.assertRaisesRegex(RuntimeError, 'mat1 dtype Double does not match mat2 dtype Float', different_dtypes)
def test_backward_noncontiguous():
# Sparse.mm backward used to wrong with non-contiguous grads,
# see https://github.com/pytorch/pytorch/issues/102493.
n_reps = 7
for _ in range(n_reps):
A = torch.eye(5).to_sparse().requires_grad_(True)
B = torch.eye(5).to_sparse()
out = torch.sparse.mm(A, B)
out.coalesce().values().sum().backward()
self.assertEqual(A.grad, A)
for n in range(2, 5):
for m in range(2, 8):
for p in range(2, 8):
test_sparse_matmul(2, 10, [n, m], [m, p])
test_sparse_matmul(2, 0, [0, 0], [0, 0])
test_sparse_matmul(2, 0, [0, 10], [10, 0])
test_error_cases()
test_backward_noncontiguous()
@coalescedonoff
@dtypes(torch.double)
def test_assign(self, device, dtype, coalesced):
def assign_to():
a, i_a, v_a = self._gen_sparse(2, 5, [2, 3], dtype, device, coalesced)
a[0] = 100
self.assertRaises(TypeError, assign_to)
@dtypes(torch.double, torch.cdouble)
def test_full_broadcast_to(self, device, dtype):
def can_broadcast(s0, s1):
s0 = tuple(reversed(s0))
s1 = tuple(reversed(s1))
for i in range(len(s0)):
if s0[i] != 1 and s0[i] != s1[i]:
return False
return True
sizes = (
(), (1,), (2,), (1, 1), (3, 1), (3, 2), (4, 1, 1), (4, 3, 2)
)
for s0, s1 in itertools.combinations(sizes, r=2):
t = make_tensor(s0, dtype=dtype, device=device, low=-9, high=9)
for sparse_dims in range(1, len(s0) + 1):
s = t.to_sparse(sparse_dims)
if can_broadcast(s0, s1):
t_res = torch.broadcast_to(t, s1)
s_res = torch._sparse_broadcast_to(s, s1)
torch._validate_sparse_coo_tensor_args(s_res._indices(), s_res._values(), s_res.shape)
if s_res.is_coalesced():
# ensure that is_coalesced is estimated correctly
self.assertEqual(s_res, torch.sparse_coo_tensor(s_res._indices(), s_res._values(), s_res.shape).coalesce())
self.assertEqual(s_res.to_dense(), t_res)
else:
with self.assertRaisesRegex(RuntimeError,
r"The expanded size of the tensor \(\d\) "
r"must match the existing size \(\d\)"):
torch._sparse_broadcast_to(s, s1)
@coalescedonoff
@dtypes(torch.double, torch.cdouble)
def test_sparse_broadcast_to(self, device, dtype, coalesced):
def test(sparse_dims, nnz, with_size, new_size):
x = self._gen_sparse(sparse_dims, nnz, with_size, dtype, device, coalesced)[0]
y = self.safeToDense(x)
x1 = torch._sparse_broadcast_to(x, new_size)
y1 = y.broadcast_to(new_size)
self.assertEqual(self.safeToDense(x1), y1)
test(4, 6, [7, 3, 1, 3, 0], [7, 3, 4, 3, 0])
test(4, 6, [7, 3, 1, 3, 0], [2, 7, 3, 1, 3, 0])
test(4, 6, [7, 3, 1, 3, 1, 3], [7, 3, 1, 3, 2, 3])
test(4, 6, [7, 3, 1, 3, 2, 1], [7, 3, 1, 3, 2, 3])
def _test_mul_skips(self, device, dtype, coalesced):
skipTestIfUncoalesced = False
# This case always coalesce inputs and that could lead to loss of precision,
# hence it is inhibited for float16/bfloat16 by providing already coalesced tensors.
if not coalesced and dtype in {torch.float16, torch.bfloat16}:
skipTestIfUncoalesced = True
# to_dense is problematic for boolean non-coalesced CUDA tensors
# see https://github.com/pytorch/pytorch/issues/81648
if not coalesced and dtype == torch.bool and torch.device(device).type == "cuda":
skipTestIfUncoalesced = True
if skipTestIfUncoalesced:
self.skipTest(f"Test with dtype={dtype}, device={device} runs only with coalesced inputs")
@coalescedonoff
# NOTE: addcmul_out is not implemented for bool.
@dtypes(*all_types_and_complex_and(torch.bfloat16, torch.float16))
@precisionOverride({torch.bfloat16: 1e-2, torch.float16: 1e-2})
def test_sparse_sparse_mul(self, device, dtype, coalesced):
self._test_mul_skips(device, dtype, coalesced)
shape = (2, 3, 4, 10)
nnz = 10
def check(self, x, y):
res_sparse = x * y
res_dense = x.to_dense() * y.to_dense()
self.assertEqual(res_sparse.to_dense(), res_dense)
def check_empty(sparse_shape, nnz, dense_shape, coalesce):
from itertools import product
for nnz_val, shape_suffix in product((nnz, 0), ((), (0,))):
empty_sparse_shape = sparse_shape + shape_suffix
empty_dense_shape = dense_shape + shape_suffix
x = self._gen_sparse(sparse_dim, nnz_val, empty_sparse_shape, dtype, device, coalesce)[0]
check(self, x, x)
# TODO: uncomment once backward is implemented for sparse tensors that broadcast in dense dims.
# def check_autograd(x, y):
# if dtype in {torch.double, torch.cdouble}:
# xa = x.detach().clone().requires_grad_(True)
# ya = y.detach().clone().requires_grad_(True)
# gradcheck(lambda a, b: (a * b).to_dense(), (xa, ya), masked=True)
# gradcheck(lambda a, b: (a * b).to_dense(), (ya, xa), masked=True)
for dim in range(len(shape) + 1):
sub_shape = shape[dim:]
sparse_dim = len(sub_shape) // 2
check_empty(sub_shape, nnz, shape, coalesced)
x = self._gen_sparse(sparse_dim, nnz, sub_shape, dtype, device, coalesced)[0]
y = self._gen_sparse(sparse_dim, nnz, sub_shape, dtype, device, coalesced)[0]
check(self, x, y)
# TODO: uncomment once supported
# check_autograd(x, y)
# check broadcasting in dense dims
for d in range(sparse_dim, len(sub_shape)):
new_shape = sub_shape[:d] + (1,) + sub_shape[d + 1:]
y = self._gen_sparse(sparse_dim, nnz, new_shape, dtype, device, coalesced)[0]
check(self, x, y)
# TODO: uncomment once supported
# check_autograd(x, y)
@coalescedonoff
@dtypes(*all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16))
@precisionOverride({torch.bfloat16: 1e-2, torch.float16: 1e-2})
def test_sparse_dense_mul(self, device, dtype, coalesced):
self._test_mul_skips(device, dtype, coalesced)
shape = (2, 3, 4, 10)
nnz = 10
def check(self, s, d):
res = d * s
# check commutativity
self.assertEqual(res, s * d)
# check correctness
self.assertEqual(res.to_dense(), s.to_dense() * d)
# check in-placeness for dense
if d.dim() >= s.dim():
dc = d.clone()
self.assertEqual(d.mul_(s), dc.mul_(s.to_dense()))
# check in-placeness for sparse
if s.dim() >= d.dim():
# for sparse
sc = s.clone()
self.assertEqual(s.mul_(d).to_dense(), sc.to_dense().mul_(d))
for dim in range(len(shape) + 1):
sub_shape = shape[dim:]
sparse_dim = len(sub_shape) // 2
def check_empty(sparse_shape, nnz, dense_shape, coalesce):
from itertools import product
for nnz_val, shape_suffix in product((nnz, 0), ((), (0,))):
empty_sparse_shape = sparse_shape + shape_suffix
empty_dense_shape = dense_shape + shape_suffix
s = self._gen_sparse(sparse_dim, nnz_val, empty_sparse_shape, dtype, device, coalesce)[0]
d = make_tensor(empty_dense_shape, dtype=dtype, device=device)
check(self, s, d)
# check scalar multiplication
s = self._gen_sparse(sparse_dim, nnz, sub_shape, dtype, device, coalesced)[0]
for scalar in (True, 1, 1.0):
res_sparse_right = s * scalar
res_sparse_left = scalar * s
res_dense = s.to_dense() * scalar
# check correctness and dtype
self.assertEqual(s.to(res_sparse_right.dtype), res_sparse_right)
self.assertEqual(res_sparse_right, res_sparse_left)
self.assertEqual(res_sparse_right.dtype, res_dense.dtype)
self.assertEqual(res_sparse_left.dtype, res_dense.dtype)
# check scalar as 0-dim sparse tensor
tscalar = torch.tensor(scalar, device=device)
sscalar = tscalar.to_sparse()
res_sparse_right = s * sscalar
res_sparse_left = sscalar * s
self.assertEqual(res_sparse_right, res_sparse_left)
self.assertEqual(s.to(res_sparse_right.dtype), res_sparse_right)
# check non-coalesced 0-dim scalar
# we skip torch.bool because for such tensors
# coalesce.to_dense != to_dense
if dtype == torch.bool:
return
for scalar_dtype in (int, float):
scalar = scalar_dtype(1)
idx = torch.tensor([], device=device).reshape(0, 2)
val = torch.tensor([scalar, scalar], device=device)
sscalar = torch.sparse_coo_tensor(idx, val, ())
res_dense = s.to_dense() * sscalar.to_dense()
self.assertEqual((s * sscalar).to_dense(), res_dense)
self.assertEqual((sscalar * s).to_dense(), res_dense)
# Case 1: sparse broadcasts over dense
s = self._gen_sparse(sparse_dim, nnz, sub_shape, dtype, device, coalesced)[0]
d = make_tensor(shape, dtype=dtype, device=device)
check(self, s, d)
check_empty(sub_shape, nnz, shape, coalesced)
# Case 2: dense broadcasts over sparse
s = self._gen_sparse(3, nnz, shape, dtype, device, coalesced)[0]
d = make_tensor(sub_shape, dtype=dtype, device=device)
check(self, s, d)
check_empty(shape, nnz, sub_shape, coalesced)
@unittest.skipIf(not TEST_NUMPY, "NumPy is not available")
@onlyCPU
@dtypes(*all_types_and_complex_and(torch.bool))
def test_sparse_spdiags(self, device, dtype):
make_diags = functools.partial(make_tensor, dtype=dtype, device=device)
make_offsets = functools.partial(torch.tensor, dtype=torch.long, device=device)
if TEST_SCIPY:
def reference(diags, offsets, shape):
return scipy.sparse.spdiags(diags, offsets, *shape).toarray()
else:
def reference(diags, offsets, shape):
result = torch.zeros(shape, dtype=dtype, device=device)
for i, off in enumerate(offsets):
res_view = result.diagonal(off)
data = diags[i]
if off > 0:
data = data[off:]
m = min(res_view.shape[0], data.shape[0])
res_view[:m] = data[:m]
return result
def check_valid(diags, offsets, shape, layout=None):
ref_out = reference(diags, offsets, shape)
out = torch.sparse.spdiags(diags, offsets, shape, layout=layout)
if layout is None:
ex_layout = torch.sparse_coo
else:
ex_layout = layout
out_dense = out.to_dense()
self.assertTrue(out.layout == ex_layout, f"Output layout {out.layout} expected {ex_layout}")
self.assertEqual(out_dense, ref_out, f"Result:\n{out_dense} does not match reference:\n{ref_out}")
def check_invalid(args, error):
with self.assertRaisesRegex(RuntimeError, error):
torch.sparse.spdiags(*args)
def valid_cases():
# some normal cases
yield (make_diags((1, 5)), make_offsets([0]), (5, 5))
yield (make_diags((3, 3)), make_offsets([-1, 0, 1]), (4, 4))
# noncontigous diags
yield (make_diags((5, 4), noncontiguous=True), make_offsets([-1, 1, 0, 2, -2]), (5, 5))
# noncontigous offsets
yield (make_diags((3, 4)), make_offsets([1, -1, 0, -2, 2])[::2], (5, 5))
# noncontigous diags + offsets
yield (make_diags((3, 4), noncontiguous=True), make_offsets([1, -1, 0, -2, 2])[::2], (5, 5))
# correct dimensionality, 2d, 2d , and shapes match, but the number of diagonals is zero
yield (make_diags((0, 3)), make_offsets([]), (3, 3))
# forward rotation of upper diagonals
yield (make_diags((3, 8)), make_offsets([1, 2, 3]), (4, 4))
# rotation exausts input space to read from
yield (make_diags((2, 3)), make_offsets([2, 1]), (3, 3))
# Simple cases repeated with special output format
yield (make_diags((1, 5)), make_offsets([0]), (5, 5), torch.sparse_csc)
yield (make_diags((3, 3)), make_offsets([-1, 0, 1]), (4, 4), torch.sparse_csr)
# vector diags
yield (make_diags((3, )), make_offsets([1]), (4, 4))
# Scalar offset
yield (make_diags((1, 3)), make_offsets(2), (4, 4))
# offsets out of range
yield (make_diags((1, 3)), make_offsets([3]), (3, 3))
yield (make_diags((1, 3)), make_offsets([-3]), (3, 3))
for case in valid_cases():
check_valid(*case)
def invalid_cases():
yield (make_diags((1, 3)), make_offsets([0]), (3, 2, 3)), "Output shape must be 2d"
yield (make_diags((2, 3)), make_offsets([[1, 2], [0, 3]]), (3, 3)), "Offsets must be scalar or vector"
yield (make_diags((3, 2, 3)), make_offsets([0, 1, 2]), (4, 4)), "Diagonals must be vector or matrix"
yield (make_diags((3, 3)), make_offsets([-1, 0]), (3, 3)), \
r"Number of diagonals \(\d\) does not match the number of offsets \(\d\)"
yield (make_diags((5,)), make_offsets([0, 1, 2, 3, 4]), (3, 3)), \
r"Number of diagonals \(\d\) does not match the number of offsets \(\d\)"
yield (make_diags((2, 2)), make_offsets([-1, 0]), (2, 3), torch.strided), \
r"Only output layouts \(\w+, \w+, \w+\) are supported, got \w+"
yield (make_diags((2, 5)), make_offsets([0, 0]), (5, 5)), "Offset tensor contains duplicate values"
yield (make_diags((1, 5)), make_offsets([0]).to(torch.int32), (5, 5)), r"Offset Tensor must have dtype Long but got \w+"
for case, error_regex in invalid_cases():
check_invalid(case, error_regex)
def test_small_nnz_coalesced(self):
# creating a coo tensor with nnz == 0 is always coalesced
self.assertTrue(torch.sparse_coo_tensor([[], []], [], (2, 2)).is_coalesced())
# same for a coo tensor with only 1 nnz
self.assertTrue(torch.sparse_coo_tensor([[0], [0]], [1], (2, 2)).is_coalesced())
# two or more nnz coalesced is false as it can't be verified without an expensive check
self.assertFalse(torch.sparse_coo_tensor([[0, 0], [0, 0]], [1, 2], (2, 2)).is_coalesced())
# even if there are no duplicates
self.assertFalse(torch.sparse_coo_tensor([[0, 1], [0, 1]], [1, 2], (2, 2)).is_coalesced())
@coalescedonoff
@dtypes(*all_types_and_complex_and(torch.bool))
def test_sum(self, device, dtype, coalesced):
def run_test(shape, nnz):
a = self._gen_sparse(2, nnz, shape, dtype, device, coalesced)[0]
self.assertEqual(a.sum(), a._values().sum())
if dtype.is_floating_point or dtype.is_complex:
a.requires_grad_(True)
a_inter = a.sum()
a_inter.abs().backward()
with torch.no_grad():
self.assertEqual(a.grad, torch.ones(shape, dtype=dtype, device=device) * torch.sgn(a_inter))
for shape in [(10, 5), (10, 10)]:
run_test(shape, 0)
run_test(shape, max(shape))
run_test(shape, shape[0] * shape[1])
class TestSparseOneOff(TestCase):
@unittest.skipIf(not TEST_CUDA, 'CUDA not available')
def test_cuda_from_cpu(self):
with self.assertRaisesRegex(
RuntimeError,
"Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!"):
torch.sparse_coo_tensor(torch.zeros(1, 4).long().cuda(),
torch.randn(4, 4, 4),
[3, 4, 4])
with self.assertRaisesRegex(
RuntimeError,
"Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!"):
torch.sparse_coo_tensor(torch.zeros(1, 4).long().cuda(),
torch.randn(4, 4, 4, 0),
[3, 4, 4, 0])
with self.assertRaisesRegex(
RuntimeError,
"Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!"):
torch.sparse_coo_tensor(torch.empty(1, 0).long().cuda(),
torch.randn(0, 4, 4, 0),
[0, 4, 4, 0])
@unittest.skipIf(not TEST_CUDA, 'CUDA not available')
def test_cuda_sparse_cpu_dense_add(self):
x = torch.zeros(3, 4, 4)
sparse_y = torch.sparse_coo_tensor(torch.zeros(1, 4).long().cuda(),
torch.randn(4, 4, 4).cuda(),
[3, 4, 4])
with self.assertRaisesRegex(RuntimeError, "add: expected 'self' to be a CUDA tensor, but got a CPU tensor"):
x + sparse_y
x = torch.zeros(3, 4, 4, 0)
sparse_y = torch.sparse_coo_tensor(torch.zeros(1, 4).long().cuda(),
torch.randn(4, 4, 4, 0).cuda(),
[3, 4, 4, 0])
with self.assertRaisesRegex(RuntimeError, "add: expected 'self' to be a CUDA tensor, but got a CPU tensor"):
x + sparse_y
x = torch.zeros(0, 4, 4, 0)
sparse_y = torch.sparse_coo_tensor(torch.empty(1, 0).long().cuda(),
torch.randn(0, 4, 4, 0).cuda(),
[0, 4, 4, 0])
with self.assertRaisesRegex(RuntimeError, "add: expected 'self' to be a CUDA tensor, but got a CPU tensor"):
x + sparse_y
def _sparse_to_dense(tensor):
if tensor.dtype != torch.bool:
return tensor.to_dense(masked_grad=True)
# to_dense uses coalesce which isn't implemented for bool
return tensor.to(torch.int8).to_dense().to(torch.bool)
_sparse_unary_ops = ops(sparse_unary_ufuncs, dtypes=OpDTypes.supported,
allowed_dtypes=all_types_and_complex())
class TestSparseUnaryUfuncs(TestCase):
exact_dtype = True
@_sparse_unary_ops
def test_sparse_consistency(self, device, dtype, op):
sample = first_sample(self, op.sample_inputs(device, dtype))
assert isinstance(sample.input, torch.Tensor)
expected = op(sample.input, *sample.args, **sample.kwargs)
assert torch.is_tensor(expected)
output = op(sample.input.to_sparse(), *sample.args, **sample.kwargs)
assert torch.is_tensor(output)
self.assertEqual(_sparse_to_dense(output), expected)
@_sparse_unary_ops
def test_out(self, device, dtype, op):
if not op.supports_out:
self.skipTest("Skipped! Out not supported")
sample = first_sample(self, op.sample_inputs(device, dtype))
sample.input = sample.input.to_sparse()
expect = op(sample.input, *sample.args, **sample.kwargs)
out = torch.sparse_coo_tensor(sample.input.shape, device=device,
dtype=expect.dtype)
op(sample.input, *sample.args, **sample.kwargs, out=out)
self.assertEqual(out, expect)
@_sparse_unary_ops
def test_inplace(self, device, dtype, op):
if op.inplace_variant is None:
self.skipTest("Skipped! Out not supported")
sample = first_sample(self, op.sample_inputs(device, dtype))
sample.input = sample.input.to_sparse().coalesce()
expect = op(sample.input, *sample.args, **sample.kwargs)
if not torch.can_cast(expect.dtype, dtype):
with self.assertRaisesRegex(RuntimeError, "result type .* can't be cast to"):
op.inplace_variant(sample.input, *sample.args, **sample.kwargs)
return
actual = op.inplace_variant(sample.input, *sample.args, **sample.kwargs)
self.assertIs(actual, sample.input)
self.assertEqual(actual, expect)
@_sparse_unary_ops
def test_sparse_zero_dims(self, device, dtype, op):
# test 0x0 sparse_coo_tensor
indices = torch.empty(2, 0, dtype=torch.int64)
values = torch.empty(0, dtype=dtype)
sparse_0x0 = torch.sparse_coo_tensor(indices, values, (0, 0))
expected = torch.sparse_coo_tensor(indices, op(values), (0, 0))
actual = op(sparse_0x0)
self.assertEqual(expected, actual)
@_sparse_unary_ops
def test_sparse_zeros(self, device, dtype, op):
samples = op.sample_inputs(device, dtype)
zero_input = torch.zeros((), device=device, dtype=dtype)
sparse_input = torch.sparse_coo_tensor((), dtype=dtype, device=device)
expect = op(zero_input)
actual = op(sparse_input)
self.assertEqual(expect, _sparse_to_dense(actual))
@ops(sparse_unary_ufuncs, dtypes=OpDTypes.supported,
allowed_dtypes=[torch.double, torch.cdouble])
def test_sparse_fn_grad(self, device, dtype, op):
if not op.supports_autograd:
self.skipTest("Skipped! Op doesn't support autograd")
for sample in op.sample_inputs(device, dtype):
sparse_input = sample.input.to_sparse().detach().requires_grad_(True)
def fn(x):
return _sparse_to_dense(
op(x, *sample.args, **sample.kwargs))
self.assertTrue(gradcheck(
fn,
(sparse_input,),
check_batched_grad=False,
check_grad_dtypes=True,
nondet_tol=op.gradcheck_nondet_tol,
fast_mode=op.gradcheck_fast_mode,
masked=True))
class TestSparseMaskedReductions(TestCase):
exact_dtype = True
fp16_low_precision_list = {
'masked.prod',
}
@ops(sparse_masked_reduction_ops)
def test_future_empty_dim(self, device, dtype, op):
"""Currently, `dim=()` in reductions operations means "reduce over
all dimensions" while in future, it will read "no reduce". See
https://github.com/pytorch/pytorch/issues/29137
For sparse masked reductions, we'll implement the current behavior.
For testing, we'll use samples with `dim=0` and map it to
`dim=()` until
torch.testing._internal.common_methods_invocations._generate_reduction_kwargs
is made to generate samples with `dim=()` for non-scalar
inputs. With this and after gh-29137 is resolved, this test
can be deleted. See also `torch.masked._canonical_dim`
implementation about changing the `dim=()` behavior.
"""
samples = op.sample_inputs_func(op, device, dtype, requires_grad=False)
op_name = op.name.replace('masked.', '')
for sample_input in samples:
if sample_input.kwargs.get('dim') != 0:
continue
sample_input_kwargs = dict(sample_input.kwargs)
sample_input_kwargs['dim'] = () # reduce over all dimensions
t = sample_input.input
mask = sample_input_kwargs.get('mask')
if mask is None and op_name in {'prod', 'amax', 'amin'}:
# FIXME: for now reductions with non-zero reduction identity and
# unspecified mask are not supported for sparse COO
# tensors, see torch.masked.prod implementation
# for details.
continue
sparse_op_kwargs = dict(sample_input_kwargs)
actual = op(t.to_sparse(), *sample_input.args, **sample_input_kwargs)
self.assertEqual(actual.layout, torch.sparse_coo)
expected = op(t, *sample_input.args, **sample_input_kwargs).to_sparse()
atol = None
rtol = None
if op.name in self.fp16_low_precision_list and dtype == torch.half:
atol = 1e-5
rtol = 2e-3
self.assertEqual(actual, expected, atol=atol, rtol=rtol)
class TestSparseMeta(TestCase):
exact_dtype = True
def _test_meta_sparse_coo(self, dtype):
r = torch.empty(4, 4, layout=torch.sparse_coo, device='meta', dtype=dtype)
self.assertTrue(r.is_meta)
self.assertEqual(r.device.type, "meta")
r2 = torch.empty_like(r)
self.assertTrue(r2.is_meta)
self.assertEqual(r, r2)
r3 = torch.sparse_coo_tensor(size=(4, 4), device='meta', dtype=dtype)
self.assertTrue(r3.is_meta)
self.assertEqual(r, r3)
r.sparse_resize_((4, 4), 1, 1)
r.sparse_resize_and_clear_((4, 4, 4), 2, 1)
self.assertEqual(r.sparse_dim(), 2)
self.assertEqual(r.dense_dim(), 1)
self.assertEqual(r._dimV(), 1)
self.assertEqual(r._nnz(), 0)
# nnz zero sparse tensors should always be coalesced at creation
self.assertEqual(r.is_coalesced(), True)
# but we can force them into the uncoalesed state
r._coalesced_(False)
self.assertEqual(r.is_coalesced(), False)
# return the coalesced state for indices/values access
r._coalesced_(True)
# TODO: this sort of aliasing will need to be handled by
# functionalization
self.assertEqual(r._indices(), torch.empty(2, 0, device='meta', dtype=torch.int64))
self.assertEqual(r._values(), torch.empty(0, 4, device='meta', dtype=dtype))
self.assertEqual(r.indices(), torch.empty(2, 0, device='meta', dtype=torch.int64))
self.assertEqual(r.values(), torch.empty(0, 4, device='meta', dtype=dtype))
def _test_meta_sparse_compressed(self, dtype, layout, batchsize, densesize):
index_dtype = torch.int64
blocksize = (2, 3) if layout in {torch.sparse_bsr, torch.sparse_bsc} else ()
sparsesize = (4, 6)
nnz = 0
shape = (*batchsize, *sparsesize, *densesize)
compressed_dim = 0 if layout in {torch.sparse_csr, torch.sparse_bsr} else 1
nof_compressed_indices = (sparsesize[compressed_dim] // blocksize[compressed_dim] + 1 if blocksize
else sparsesize[compressed_dim] + 1)
compressed_indices = torch.empty((*batchsize, nof_compressed_indices), device='meta', dtype=index_dtype)
plain_indices = torch.empty((*batchsize, nnz), device='meta', dtype=index_dtype)
values = torch.empty((*batchsize, nnz, *blocksize, *densesize), device='meta', dtype=dtype)
r = torch.sparse_compressed_tensor(
compressed_indices,
plain_indices,
values,
shape,
layout=layout
)
self.assertTrue(r.is_meta)
self.assertEqual(r.device.type, "meta")
self.assertEqual(r.sparse_dim(), 2)
self.assertEqual(r.dense_dim(), len(densesize))
self.assertEqual(r._nnz(), nnz)
batch_dims = r.ndim - r.sparse_dim() - r.dense_dim()
r_blocksize = r.values().shape[batch_dims + 1: batch_dims + 1 + len(blocksize)]
self.assertEqual(r_blocksize, blocksize)
r_compressed_indices = r.crow_indices() if layout in {torch.sparse_csr, torch.sparse_bsr} else r.ccol_indices()
r_plain_indices = r.col_indices() if layout in {torch.sparse_csr, torch.sparse_bsr} else r.row_indices()
self.assertEqual(r_compressed_indices,
torch.empty((*batchsize, nof_compressed_indices), device='meta', dtype=index_dtype))
self.assertEqual(r_plain_indices, torch.empty((*batchsize, nnz), device='meta', dtype=index_dtype))
self.assertEqual(r.values(), torch.empty((*batchsize, nnz, *blocksize, *densesize), device='meta', dtype=dtype))
r2 = torch.empty_like(r)
self.assertTrue(r2.is_meta)
self.assertEqual(r2, r)
if layout in {torch.sparse_csr, torch.sparse_csc}:
r3 = torch.empty((*batchsize, *sparsesize), dtype=dtype, layout=layout, device="meta")
self.assertTrue(r3.is_meta)
if not densesize:
# dense dimensions cannot be specified for torch.empty
self.assertEqual(r3, r)
@all_sparse_layouts('layout', include_strided=False)
@parametrize("dtype", [torch.float64])
def test_meta(self, dtype, layout):
if layout is torch.sparse_coo:
self._test_meta_sparse_coo(dtype)
else:
for batchsize, densesize in itertools.product([(), (2,)], [(), (3,)]):
self._test_meta_sparse_compressed(dtype, layout, batchsize, densesize)
def _test_print_meta_data(self, dtype, layout, batchsize, sparsesize, densesize):
index_dtype = torch.int64
nnz = 0
blocksize = (2, 3) if layout in {torch.sparse_bsr, torch.sparse_bsc} else ()
shape = (*batchsize, *sparsesize, *densesize)
values = torch.empty((*batchsize, nnz, *blocksize, *densesize), device='meta', dtype=dtype)
if layout is torch.sparse_coo:
indices = torch.empty((len(sparsesize), nnz), device='meta', dtype=index_dtype)
x = torch.sparse_coo_tensor(indices, values, shape)
else:
compressed_dim = 0 if layout in {torch.sparse_csr, torch.sparse_bsr} else 1
nof_compressed_indices = (sparsesize[compressed_dim] // blocksize[compressed_dim] + 1 if blocksize
else sparsesize[compressed_dim] + 1)
compressed_indices = torch.empty((*batchsize, nof_compressed_indices), device='meta', dtype=index_dtype)
plain_indices = torch.empty((*batchsize, nnz), device='meta', dtype=index_dtype)
x = torch.sparse_compressed_tensor(
compressed_indices,
plain_indices,
values,
shape,
layout=layout
)
printed = []
printed.append(f"########## {dtype}/{index_dtype}/size={batchsize}+{sparsesize}+{blocksize}+{densesize} ##########")
printed.append("# sparse meta tensor")
printed.append(str(x))
return printed
@all_sparse_layouts('layout', include_strided=False)
@parametrize("dtype", [torch.float64])
def test_print_meta(self, dtype, layout):
printed = []
for batchsize, sparsesize, densesize in itertools.product(
[(), (2,)], [(4, 6), (3, 5, 7)], [(), (3,)]
):
if layout is torch.sparse_coo and batchsize:
# COO tensors don't have batch dimensions
continue
if layout is not torch.sparse_coo and len(sparsesize) != 2:
# CSR/CSC/BSR/BSC tensors must have 2 sparse dimensions
continue
printed += self._test_print_meta_data(dtype, layout, batchsize, sparsesize, densesize)
orig_maxDiff = self.maxDiff
self.maxDiff = None
try:
self.assertExpected('\n'.join(printed))
self.maxDiff = orig_maxDiff
except Exception:
self.maxDiff = orig_maxDiff
raise
def assertEqualMeta(self, x, y, expected_nnz):
self.assertEqual(x.layout, y.layout)
self.assertEqual(x.shape, y.shape)
self.assertEqual(x.dtype, y.dtype)
self.assertEqual(x.sparse_dim(), y.sparse_dim())
self.assertEqual(x.dense_dim(), y.dense_dim())
def assertEqualAttrs(x, y, expected_shape):
self.assertEqual(x.shape, expected_shape)
self.assertEqual(x.dtype, y.dtype)
self.assertEqual(x.layout, y.layout)
if not x.is_meta:
self.assertEqual(x.device, y.device)
if x.layout is torch.sparse_coo:
assertEqualAttrs(x._indices(), y._indices(), (*y._indices().shape[:-1], expected_nnz))
assertEqualAttrs(x._values(), y._values(), (expected_nnz, *y._values().shape[1:]))
elif x.layout in {torch.sparse_csr, torch.sparse_bsr}:
assertEqualAttrs(x.crow_indices(), y.crow_indices(), y.crow_indices().shape)
assertEqualAttrs(x.col_indices(), y.col_indices(), (*y.col_indices().shape[:-1], expected_nnz))
batch_dim = x.col_indices().ndim - 1
values_shape = (*y.values().shape[:batch_dim], expected_nnz, *y.values().shape[batch_dim + 1:])
self.assertEqual(x.values().layout, y.values().layout)
self.assertEqual(x.values().dtype, y.values().dtype)
self.assertEqual(x.values().shape, values_shape)
elif x.layout in {torch.sparse_csc, torch.sparse_bsc}:
assertEqualAttrs(x.ccol_indices(), y.ccol_indices(), y.ccol_indices().shape)
assertEqualAttrs(x.row_indices(), y.row_indices(), (*y.row_indices().shape[:-1], expected_nnz))
batch_dim = x.row_indices().ndim - 1
values_shape = (*y.values().shape[:batch_dim], expected_nnz, *y.values().shape[batch_dim + 1:])
self.assertEqual(x.values().layout, y.values().layout)
self.assertEqual(x.values().dtype, y.values().dtype)
self.assertEqual(x.values().shape, values_shape)
@all_sparse_layouts('layout', include_strided=False)
@parametrize("dtype", [torch.float64])
def test_to_meta(self, dtype, layout):
index_dtype = torch.int64
device = 'cpu'
for t in self.generate_simple_inputs(layout, device=device, dtype=dtype, index_dtype=index_dtype):
m = t.to(device="meta")
self.assertEqual(m.device.type, "meta")
self.assertEqualMeta(m, t, 0)
@all_sparse_layouts('layout', include_strided=False)
@parametrize("dtype", [torch.float64])
def test_zeros_like_meta(self, dtype, layout):
index_dtype = torch.int64
device = 'cpu'
for t in self.generate_simple_inputs(layout, device=device, dtype=dtype, index_dtype=index_dtype):
m = torch.zeros_like(t, device="meta")
self.assertEqual(m.device.type, "meta")
self.assertEqualMeta(m, t, 0)
@all_sparse_layouts('layout', include_strided=False)
@parametrize("dtype", [torch.float64])
def test_fake(self, dtype, layout):
from torch._subclasses.fake_tensor import FakeTensorMode, FakeTensor
fake_mode = FakeTensorMode()
index_dtype = torch.int64
device = 'cpu'
for t in self.generate_simple_inputs(layout, device=device, dtype=dtype, index_dtype=index_dtype):
f = FakeTensor.from_tensor(t, fake_mode)
self.assertIsInstance(f, FakeTensor)
self.assertEqualMeta(f, t, 0)
d = f.detach()
self.assertIsInstance(d, FakeTensor)
self.assertEqualMeta(d, t, 0)
@all_sparse_layouts('layout', include_strided=False)
@parametrize("dtype", [torch.float64])
def test_zeros_like_fake(self, dtype, layout):
from torch._subclasses.fake_tensor import FakeTensorMode, FakeTensor
from torch.utils._mode_utils import no_dispatch
fake_mode = FakeTensorMode()
index_dtype = torch.int64
device = 'cpu'
for t in self.generate_simple_inputs(layout, device=device, dtype=dtype, index_dtype=index_dtype):
f = FakeTensor.from_tensor(t, fake_mode)
expected = torch.zeros_like(t)
with no_dispatch():
result = torch.zeros_like(f, device=f.fake_device)
self.assertEqual(result, expected)
self.assertEqualMeta(result, expected, 0)
@all_sparse_layouts('layout', include_strided=False)
@parametrize("dtype", [torch.float64])
def test_sum_meta(self, dtype, layout):
device = 'cpu'
index_dtype = torch.int64
for t in self.generate_simple_inputs(layout, device=device, dtype=dtype, index_dtype=index_dtype):
m = t.to(device='meta')
r = torch.sum(m)
expected = torch.sum(t).to(device="meta")
self.assertTrue(r.is_meta)
self.assertEqualMeta(r, expected, 0)
@all_sparse_layouts('layout', include_strided=False)
@parametrize("dtype", [torch.float64])
def test_add_meta(self, dtype, layout):
device = 'cpu'
index_dtype = torch.int64
for t in self.generate_simple_inputs(layout, device=device, dtype=dtype, index_dtype=index_dtype):
expected = torch.add(t, t).to(device='meta')
m = t.to(device='meta')
r = torch.add(m, m)
self.assertEqualMeta(r, expected, 0)
class _SparseDataset(torch.utils.data.Dataset):
# An utility class used in TestSparseAny.test_dataloader method.
def __init__(self, sparse_tensors):
self.sparse_tensors = sparse_tensors
def __len__(self):
return len(self.sparse_tensors)
def __getitem__(self, index):
return self.sparse_tensors[index]
class TestSparseAny(TestCase):
@onlyCPU
@all_sparse_layouts('layout', include_strided=False)
@torch.sparse.check_sparse_tensor_invariants(enable=False)
def test_check_sparse_tensor_invariants(self, layout):
if layout is torch.sparse_coo:
def create_invalid_tensor(check_invariants=None):
shape = (2, 2)
invalid_indices = torch.tensor([[0], [3]]) # column index is out of range
values = torch.tensor([1])
if check_invariants is None:
return torch.sparse_coo_tensor(invalid_indices, values, shape)
else:
return torch.sparse_coo_tensor(invalid_indices, values, shape, check_invariants=check_invariants)
expected_exception_message = 'size is inconsistent with indices: for dim 1, size is 2 but found index 3'
elif layout in {torch.sparse_csr, torch.sparse_csc, torch.sparse_bsr, torch.sparse_bsc}:
def create_invalid_tensor(check_invariants=None):
shape = (2, 2)
compressed_indices = torch.tensor([0, 0, 1])
invalid_plain_indices = torch.tensor([3]) # index is out of range
if layout in {torch.sparse_bsr, torch.sparse_bsc}:
values = torch.tensor([[[1]]])
else:
values = torch.tensor([1])
if check_invariants is None:
return torch.sparse_compressed_tensor(compressed_indices, invalid_plain_indices, values, shape, layout=layout)
else:
return torch.sparse_compressed_tensor(compressed_indices, invalid_plain_indices, values, shape, layout=layout,
check_invariants=check_invariants)
if layout in {torch.sparse_csr, torch.sparse_bsr}:
expected_exception_message = r'`0 <= col_indices < ncols` is not satisfied.'
else:
expected_exception_message = r'`0 <= row_indices < nrows` is not satisfied.'
else:
raise NotImplementedError(layout)
# First, consider the case where invariant checks are disabled
# "globally" (read: within the context of this test method
# caller) as defined by check_sparse_tensor_invariants(False)
# decorator:
self.assertFalse(torch.sparse.check_sparse_tensor_invariants.is_enabled())
# Enable the invariant checks in a local context:
with torch.sparse.check_sparse_tensor_invariants():
self.assertTrue(torch.sparse.check_sparse_tensor_invariants.is_enabled())
# Leaving the local context must restore the "global" state of
# the invariant check feature:
self.assertFalse(torch.sparse.check_sparse_tensor_invariants.is_enabled())
# Since invariant checks are disabled by default, we can
# create an invalid sparse tensor without raising an
# exception:
r = create_invalid_tensor()
self.assertEqual(r.layout, layout)
# Or, when disabling the invariants check explicitly:
r = create_invalid_tensor(check_invariants=False)
self.assertEqual(r.layout, layout)
# Enabling invariant check via constructor's optional argument
# will raise an exception when sparse tensor invariants are
# violated:
with self.assertRaisesRegex(RuntimeError, expected_exception_message):
create_invalid_tensor(check_invariants=True)
# Check that the global invariant check flag has been restored
# after raising the exception above:
self.assertFalse(torch.sparse.check_sparse_tensor_invariants.is_enabled())
# Next, consider the case where invariant checks are enabled
# within a local context:
with torch.sparse.check_sparse_tensor_invariants():
self.assertTrue(torch.sparse.check_sparse_tensor_invariants.is_enabled())
# Since invariant checks are now enabled by default, an
# attempt to create an invalid sparse tensor will lead to
# an exception:
with self.assertRaisesRegex(RuntimeError, expected_exception_message):
create_invalid_tensor()
# Similarly, when enabling the invariant checks
# explicitly, invalid sparse tensor construction will lead
# to an exception:
with self.assertRaisesRegex(RuntimeError, expected_exception_message):
create_invalid_tensor(check_invariants=True)
# However, invariants check can be disabled via
# constructor's optional argument so that the invalid
# tensor is succesfully constructed:
r = create_invalid_tensor(check_invariants=False)
self.assertEqual(r.layout, layout)
# Check that the invariant check flag has been restored
# when leaving the constructor:
self.assertTrue(torch.sparse.check_sparse_tensor_invariants.is_enabled())
# Double-check restoring the global state when leaving the
# local context:
self.assertFalse(torch.sparse.check_sparse_tensor_invariants.is_enabled())
# Test nesting of pre-defined context managers
check_ctx = torch.sparse.check_sparse_tensor_invariants(True)
no_check_ctx = torch.sparse.check_sparse_tensor_invariants(False)
with check_ctx:
self.assertTrue(torch.sparse.check_sparse_tensor_invariants.is_enabled())
with no_check_ctx:
self.assertFalse(torch.sparse.check_sparse_tensor_invariants.is_enabled())
self.assertTrue(torch.sparse.check_sparse_tensor_invariants.is_enabled())
self.assertFalse(torch.sparse.check_sparse_tensor_invariants.is_enabled())
# Test an attempt to re-use an activate context manager instance
check_ctx2 = torch.sparse.check_sparse_tensor_invariants(True)
with check_ctx:
self.assertTrue(torch.sparse.check_sparse_tensor_invariants.is_enabled())
with no_check_ctx:
self.assertFalse(torch.sparse.check_sparse_tensor_invariants.is_enabled())
with self.assertRaisesRegex(RuntimeError, "This context manager instance is already activated."
" Use a different context manager instance for context nesting"):
with check_ctx:
self.assertTrue(torch.sparse.check_sparse_tensor_invariants.is_enabled())
self.assertFalse(torch.sparse.check_sparse_tensor_invariants.is_enabled())
with check_ctx2:
self.assertTrue(torch.sparse.check_sparse_tensor_invariants.is_enabled())
self.assertFalse(torch.sparse.check_sparse_tensor_invariants.is_enabled())
self.assertTrue(torch.sparse.check_sparse_tensor_invariants.is_enabled())
self.assertFalse(torch.sparse.check_sparse_tensor_invariants.is_enabled())
def test_generate_simple_inputs(self):
layouts = [torch.strided, torch.sparse_coo, torch.sparse_csr, torch.sparse_csc, torch.sparse_bsr, torch.sparse_bsc]
tested_combinations = set()
for tensors in zip(*map(self.generate_simple_inputs, layouts)):
for i, t in enumerate(tensors):
self.assertEqual(t.layout, layouts[i])
# all layouts must produce semantically the same tensors
self.assertEqual(t, tensors[0])
if t.layout is torch.strided:
is_hybrid = None
else:
is_hybrid = t.dense_dim() > 0
if t.layout in {torch.sparse_csr, torch.sparse_bsr}:
is_batch = t.crow_indices().ndim > 1
elif t.layout in {torch.sparse_csc, torch.sparse_bsc}:
is_batch = t.ccol_indices().ndim > 1
else:
is_batch = None
if t.layout in {torch.sparse_bsr, torch.sparse_bsc}:
blocksize = t.values().shape[1:3]
nontrivial_blocksize = 1 not in blocksize
else:
nontrivial_blocksize = None
if t.layout in {torch.sparse_csr, torch.sparse_bsr}:
contiguous_indices = t.crow_indices().is_contiguous() and t.col_indices().is_contiguous()
contiguous_values = t.values().is_contiguous()
elif t.layout in {torch.sparse_csc, torch.sparse_bsc}:
contiguous_indices = t.ccol_indices().is_contiguous() and t.row_indices().is_contiguous()
contiguous_values = t.values().is_contiguous()
elif t.layout is torch.sparse_coo:
contiguous_indices = t._indices().is_contiguous()
contiguous_values = t._values().is_contiguous()
else:
contiguous_indices = None
contiguous_values = t.is_contiguous()
tested_combinations.add((t.layout, is_hybrid, is_batch, nontrivial_blocksize,
contiguous_indices, contiguous_values))
# Ensure that the inputs generation covers all layout,
# non-hybrid/hybrid, non-batch/batch, and contiguity
# combinations:
untested_combinations = set()
for layout in layouts:
for is_hybrid in [False, True]:
if layout is torch.strided:
is_hybrid = None
for is_batch in [False, True]:
if layout in {torch.sparse_coo, torch.strided}:
is_batch = None
for nontrivial_blocksize in [False, True]:
if layout not in {torch.sparse_bsr, torch.sparse_bsc}:
nontrivial_blocksize = None
for contiguous_indices in [False, True]:
if layout is torch.strided:
contiguous_indices = None
elif not is_batch:
# indices are contiguous per-patch
contiguous_indices = True
for contiguous_values in [False, True]:
key = (layout, is_hybrid, is_batch, nontrivial_blocksize,
contiguous_indices, contiguous_values)
if key not in tested_combinations:
untested_combinations.add(
f'layout={layout}, is_hybrid={is_hybrid}, is_batch={is_batch},'
f' nontrivial_blocksize={nontrivial_blocksize},'
f' contiguous_indices{contiguous_indices}, contiguous_values={contiguous_values}')
assert not untested_combinations, untested_combinations
@all_sparse_layouts('layout', include_strided=False)
def test_constructor_autograd(self, device, layout):
def specific_constructor(*args, **kwargs):
if layout is torch.sparse_csr:
return torch.sparse_csr_tensor(*args, **kwargs)
elif layout is torch.sparse_csc:
return torch.sparse_csc_tensor(*args, **kwargs)
elif layout is torch.sparse_bsc:
return torch.sparse_bsc_tensor(*args, **kwargs)
elif layout is torch.sparse_bsr:
return torch.sparse_bsr_tensor(*args, **kwargs)
elif layout is torch.sparse_coo:
return torch.sparse_coo_tensor(*args, **kwargs)
else:
raise NotImplementedError(layout)
def generic_constructor(*args, **kwargs):
if layout in {torch.sparse_csr, torch.sparse_csc, torch.sparse_bsr, torch.sparse_bsc}:
kwargs.update(layout=layout)
return torch.sparse_compressed_tensor(*args, **kwargs)
elif layout is torch.sparse_coo:
return torch.sparse_coo_tensor(*args, **kwargs)
else:
raise NotImplementedError(layout)
if layout is torch.sparse_coo:
constructors = (specific_constructor,)
else:
constructors = (specific_constructor, generic_constructor)
for args, kwargs in self.generate_simple_inputs(
layout, device=device, dtype=torch.float64,
enable_batch=False, # TODO: remove after gh-104868 is resolved
output_tensor=False):
values_offset = 1 if layout is torch.sparse_coo else 2
for cnstr in constructors:
for requires_grad in (False, True):
values = args[values_offset].detach().requires_grad_(requires_grad)
args = (*args[:values_offset], values, *args[values_offset + 1:])
kwargs_ = dict(kwargs)
args_ = args + (kwargs_.pop('size'),)
sparse = cnstr(*args, **kwargs)
self.assertEqual(sparse.requires_grad, requires_grad)
if requires_grad:
for masked in (False, True):
if layout is torch.sparse_coo:
torch.autograd.gradcheck(
lambda i, v: cnstr(i, v, **kwargs).to_dense(masked_grad=masked),
args, masked=masked)
torch.autograd.gradcheck(
lambda i, v, sz: cnstr(i, v, sz, **kwargs_).to_dense(masked_grad=masked),
args_, masked=masked)
else:
if layout in {torch.sparse_csc, torch.sparse_bsr, torch.sparse_bsc} and 0:
# TODO: remove this if-block after gh-107370 is resolved
continue
torch.autograd.gradcheck(
lambda ci, pi, v: cnstr(ci, pi, v, **kwargs).to_dense(masked_grad=masked),
args, masked=masked)
torch.autograd.gradcheck(
lambda ci, pi, v, sz: cnstr(ci, pi, v, sz, **kwargs_).to_dense(masked_grad=masked),
args_, masked=masked)
@all_sparse_layouts('from_layout', include_strided=False)
@dtypes(*all_types_and_complex_and(torch.half, torch.bool, torch.bfloat16))
@parametrize("index_dtype", [torch.int32, torch.int64])
def test_to_dense(self, from_layout, device, dtype, index_dtype):
"""
This test tests conversion from any layout to strided layout.
"""
for t in self.generate_simple_inputs(
from_layout, device=device, dtype=dtype, index_dtype=index_dtype):
r = t.to_dense()
self.assertEqual(r.layout, torch.strided)
self.assertEqual(r, t)
@all_sparse_layouts('from_layout', include_strided=False)
@dtypes(torch.float64, torch.complex128)
@parametrize("index_dtype", [torch.int64])
@gradcheck_semantics()
def test_gradcheck_to_dense(self, from_layout, device, dtype, index_dtype, gradcheck):
for t in self.generate_simple_inputs(
from_layout, device=device, dtype=dtype, index_dtype=index_dtype):
batch_dim = t.dim() - t.dense_dim() - t.sparse_dim()
if batch_dim > 0:
# TODO: implement batch support in _convert_indices_from_csr_to_coo
continue
t = t.clone().detach().requires_grad_(True)
r = gradcheck(lambda x: torch.Tensor.to_dense(x, masked_grad=gradcheck.masked), t)
self.assertTrue(r)
@all_sparse_layouts('from_layout', include_strided=True)
@all_sparse_layouts('to_layout', include_strided=False)
@dtypes(*all_types_and_complex_and(torch.half, torch.bool, torch.bfloat16))
@parametrize("index_dtype", [torch.int32, torch.int64])
def test_to_sparse(self, from_layout, to_layout, device, dtype, index_dtype):
"""
This test tests conversion from any layout to any sparse layout.
"""
for t in self.generate_simple_inputs(
from_layout, device=device, dtype=dtype, index_dtype=index_dtype,
enable_hybrid=(
# TODO: to support conversion strided->hybrid
# CSR/CSC/BSR/BSC, to_sparse() requires extra keyword
# argument, either nof_batch_dims or
# nof_dense_dims
not (from_layout is torch.strided and to_layout in
{torch.sparse_bsr, torch.sparse_bsc, torch.sparse_csr, torch.sparse_csc}))):
if to_layout in {torch.sparse_bsr, torch.sparse_bsc}:
if from_layout == torch.sparse_bsr:
batch_ndim = t.crow_indices().dim() - 1
blocksize = t.values().shape[batch_ndim + 1:batch_ndim + 3]
elif from_layout == torch.sparse_bsc:
batch_ndim = t.ccol_indices().dim() - 1
blocksize = t.values().shape[batch_ndim + 1:batch_ndim + 3]
else:
blocksize = (1, 1)
else:
blocksize = None
if from_layout is torch.strided:
is_batch = None
is_hybrid = None
else:
is_batch = t.dim() > (t.sparse_dim() + t.dense_dim())
is_hybrid = t.dense_dim() > 0
def explicit_to_sparse(x):
# Used to check that the explicit conversion methods
# are consistent with the `to_sparse(*, layout,
# blocksize)` method.
if to_layout is torch.sparse_coo:
return x.to_sparse_coo()
elif to_layout is torch.sparse_csr:
return x.to_sparse_csr()
elif to_layout is torch.sparse_csc:
return x.to_sparse_csc()
elif to_layout is torch.sparse_bsr:
return x.to_sparse_bsr(blocksize)
elif to_layout is torch.sparse_bsc:
return x.to_sparse_bsc(blocksize)
else:
assert 0 # unreachable
# TODO: The following exception cases all correspond to
# not implemented conversions
if from_layout in {
torch.sparse_csr, torch.sparse_csc} and to_layout in {torch.sparse_bsr, torch.sparse_bsc} and is_batch:
with self.assertRaisesRegex(
RuntimeError,
r"conversion from Sparse(Csr|Csc) to Sparse(Bsr|Bsc) for batched inputs is not supported"):
t.to_sparse(layout=to_layout, blocksize=blocksize)
with self.assertRaisesRegex(
RuntimeError,
r"conversion from Sparse(Csr|Csc) to Sparse(Bsr|Bsc) for batched inputs is not supported"):
explicit_to_sparse(t)
continue
elif from_layout is torch.sparse_coo and to_layout in {
torch.sparse_csr, torch.sparse_csc, torch.sparse_bsr, torch.sparse_bsc} and t.sparse_dim() != 2:
with self.assertRaisesRegex(
RuntimeError,
r"conversion from Sparse to .* for input tensors with sparse_dim\(\)!=2 is not supported"):
t.to_sparse(layout=to_layout, blocksize=blocksize)
with self.assertRaisesRegex(
RuntimeError,
r"conversion from Sparse to .* for input tensors with sparse_dim\(\)!=2 is not supported"):
explicit_to_sparse(t)
continue
elif (from_layout, to_layout) in {(torch.sparse_bsc, torch.sparse_csr), (torch.sparse_bsc, torch.sparse_csc),
(torch.sparse_bsr, torch.sparse_csr), (torch.sparse_bsr, torch.sparse_csc)}:
with self.assertRaisesRegex(
RuntimeError,
r"sparse_compressed_to_sparse_(csr|csc|bsr|bsc): expected\s*(Sparse(Csc|Csr)[,]|)\s*Sparse(Csr|Bsr)"
" or Sparse(Csc|Bsc) layout but got Sparse(Csr|Csc|Bsr|Bsc)"):
t.to_sparse(layout=to_layout, blocksize=blocksize)
with self.assertRaisesRegex(
RuntimeError,
r"sparse_compressed_to_sparse_(csr|csc|bsr|bsc): expected\s*(Sparse(Csc|Csr)[,]|)\s*Sparse(Csr|Bsr)"
" or Sparse(Csc|Bsc) layout but got Sparse(Csr|Csc|Bsr|Bsc)"):
explicit_to_sparse(t)
self.skipTest('NOT IMPL')
else:
r = t.to_sparse(layout=to_layout, blocksize=blocksize)
self.assertEqual(r.layout, to_layout)
# to_sparse method uses unsafe construction of sparse
# tensors. Here we explicitly validate the results to
# make sure that the sparse tensors are consistent
# with the corresponding sparse tensor invariants.
if r.layout in {torch.sparse_csr, torch.sparse_bsr, torch.sparse_csc, torch.sparse_bsc}:
if r.layout in {torch.sparse_csr, torch.sparse_bsr}:
compressed_indices, plain_indices = r.crow_indices(), r.col_indices()
else:
compressed_indices, plain_indices = r.ccol_indices(), r.row_indices()
torch._validate_sparse_compressed_tensor_args(compressed_indices, plain_indices, r.values(),
r.shape, r.layout)
if from_layout in {torch.strided, torch.sparse_coo}:
self.assertEqual(compressed_indices.dtype, torch.int64)
self.assertEqual(plain_indices.dtype, torch.int64)
else:
self.assertEqual(compressed_indices.dtype, index_dtype)
self.assertEqual(plain_indices.dtype, index_dtype)
self.assertEqual(r.values().dtype, dtype)
elif r.layout is torch.sparse_coo:
if t.layout is torch.sparse_coo:
self.assertEqual(t.is_coalesced(), r.is_coalesced())
# Check r is truly coalesced when r.is_coalesced == True
if r.is_coalesced():
self.assertTrue(is_coalesced_indices(r))
torch._validate_sparse_coo_tensor_args(r._indices(), r._values(), r.shape)
self.assertEqual(r._indices().dtype, torch.int64)
self.assertEqual(r._values().dtype, dtype)
else:
assert 0 # unreachable
# Finally, we'll test tensor equality:
self.assertEqual(r, t)
# Also, check consistency with explicit conversion methods:
r2 = explicit_to_sparse(t)
self.assertEqual(r2, r)
# Check inverse conversion from sparse compressed block tensors
if from_layout == torch.sparse_bsr:
batch_ndim = t.crow_indices().dim() - 1
from_blocksize = t.values().shape[batch_ndim + 1:batch_ndim + 3]
elif from_layout == torch.sparse_bsc:
batch_ndim = t.ccol_indices().dim() - 1
from_blocksize = t.values().shape[batch_ndim + 1:batch_ndim + 3]
else:
continue
if r.ndim != 2:
continue
t2 = r.to_sparse(layout=from_layout, blocksize=from_blocksize)
self.assertEqual(t2, t)
# extra tests
if (from_layout, to_layout) == (torch.sparse_csr, torch.sparse_bsr):
# See gh-90910
t = torch.tensor([[0, 0, 1, 0], [0, 1, 0, 0]], dtype=dtype, device=device).to_sparse_csr()
r = t.to_sparse_bsr((2, 2))
torch._validate_sparse_compressed_tensor_args(r.crow_indices(), r.col_indices(), r.values(), r.shape, r.layout)
self.assertEqual(r, t)
if (from_layout, to_layout) in {(torch.sparse_csr, torch.sparse_csc),
(torch.sparse_csc, torch.sparse_csr)}:
# See gh-91007
compressed_indices = torch.tensor([0, 4, 8, 8, 12, 16, 20], dtype=index_dtype, device=device)
plain_indices = torch.tensor([0, 1, 2, 3] * 5, dtype=index_dtype, device=device)
t = torch.sparse_compressed_tensor(compressed_indices, plain_indices, range(20),
dtype=dtype, device=device, layout=from_layout)
r = t.to_sparse(layout=to_layout)
if r.layout in {torch.sparse_csr, torch.sparse_bsr}:
compressed_indices, plain_indices = r.crow_indices(), r.col_indices()
else:
compressed_indices, plain_indices = r.ccol_indices(), r.row_indices()
torch._validate_sparse_compressed_tensor_args(compressed_indices, plain_indices, r.values(), r.shape, r.layout)
self.assertEqual(r, t)
@onlyNativeDeviceTypes
@suppress_warnings
@ops(reduction_ops_with_sparse_support)
@precisionOverride({torch.bfloat16: 5e-4, torch.float16: 5e-3})
@all_sparse_layouts('layout', include_strided=False)
def test_reductions(self, layout, device, dtype, op):
count = 0
for sample in op.sample_inputs_sparse(layout, device, dtype):
count += 1
t_inp, t_args, t_kwargs = sample.input, sample.args, sample.kwargs
result = op.op(t_inp, *t_args, **t_kwargs)
# Checking invariant rop(inp, ...).to_dense() == rop(inp.to_dense(), ...)
dense = op.op(t_inp.to_dense(), *t_args, **t_kwargs)
self.assertEqual(result, dense)
if count == 0:
# we count samples to avoid false-positive test reports
self.skipTest('no sample inputs')
@onlyNativeDeviceTypes
@suppress_warnings
@ops(reduction_ops_with_sparse_support, allowed_dtypes=(torch.float32, torch.float64, torch.complex64, torch.complex128))
@all_sparse_layouts('layout', include_strided=False)
def test_reductions_backward(self, layout, device, dtype, op):
count = 0
for sample in op.sample_inputs_sparse(layout, device, dtype, requires_grad=True):
t_inp, t_args, t_kwargs = sample.input, sample.args, sample.kwargs
r = op.op(t_inp, *t_args, **t_kwargs)
if r.numel() != 0:
r = r.sum()
if op.name == 'sum':
count += 1
r.abs().backward()
self.assertEqual(t_inp.grad, torch.ones(t_inp.shape, dtype=dtype, device=device) * torch.sgn(r))
else:
self.skipTest('NOT IMPL')
if count == 0:
# we count samples to avoid false-positive test reports
self.skipTest('no sample inputs')
@onlyNativeDeviceTypes
@suppress_warnings
@parametrize("mth", [subtest(mth, name=mth.__name__)
for mth in [torch.Tensor.is_coalesced,
torch.Tensor.coalesce,
torch.Tensor.indices,
torch.Tensor.values,
torch.Tensor.crow_indices,
torch.Tensor.col_indices,
torch.Tensor.ccol_indices,
torch.Tensor.row_indices,
]])
@all_sparse_layouts('layout', include_strided=True)
def test_unsupported_backend_error_message(self, mth, layout, device):
inp = torch.tensor([[1, 2], [3, 4]], device=device).to_sparse(
layout=layout,
blocksize=(1, 1) if layout in {torch.sparse_bsr, torch.sparse_bsc} else None)
assert inp.layout is layout
expected_behaviour = dict(
# <mth name> = (<supported layouts>, <exception message on other layouts>)
is_coalesced=({torch.sparse_coo},
"is_coalesced expected sparse coordinate tensor layout but got (Sparse(Csr|Csc|Bsr|Bsc)|Strided)"),
coalesce=({torch.sparse_coo},
"coalesce expected sparse coordinate tensor layout but got (Sparse(Csr|Csc|Bsr|Bsc)|Strided)"),
indices=({torch.sparse_coo},
"indices expected sparse coordinate tensor layout but got (Sparse(Csr|Csc|Bsr|Bsc)|Strided)"),
values=({torch.sparse_coo, torch.sparse_csr, torch.sparse_csc, torch.sparse_bsr, torch.sparse_bsc},
"values expected sparse tensor layout but got Strided"),
crow_indices=({torch.sparse_csr, torch.sparse_bsr},
"crow_indices expected sparse row compressed tensor layout but got (Sparse(Csc|Bsc|)|Strided)"),
col_indices=({torch.sparse_csr, torch.sparse_bsr},
"col_indices expected sparse row compressed tensor layout but got (Sparse(Csc|Bsc|)|Strided)"),
ccol_indices=({torch.sparse_csc, torch.sparse_bsc},
"ccol_indices expected sparse column compressed tensor layout but got (Sparse(Csr|Bsr|)|Strided)"),
row_indices=({torch.sparse_csc, torch.sparse_bsc},
"row_indices expected sparse column compressed tensor layout but got (Sparse(Csr|Bsr|)|Strided)"),
)[mth.__name__]
if layout in expected_behaviour[0]:
mth(inp)
else:
with self.assertRaisesRegex(RuntimeError, expected_behaviour[1]):
mth(inp)
@onlyNativeDeviceTypes
@all_sparse_layouts('layout', include_strided=not True)
@dtypes(torch.float64, torch.cdouble)
@parametrize("masked", [subtest(False, name='sparse'), subtest(True, name='masked')])
@parametrize("fast_mode", [subtest(False, name='slow'), subtest(True, name='fast')])
def test_gradcheck_mm(self, layout, dtype, device, masked, fast_mode):
# This function does not check the following cases:
# - batch or hybrid tensors because addmm does not support
# such inputs yet
# - check_forward_ad=True because of the lack of sparse tensor
# support in aten::view_as_real, torch._VF._make_dual, etc.
ref_x = torch.tensor([[1, 2, 0, 0],
[0, 6, 0, 0],
[0, 0, 0, 0],
[13, 14, 0, 15]], dtype=dtype, device=device)
ref_y = torch.tensor([[11, 12, 13, 14],
[21, 22, 23, 24],
[31, 32, 33, 34],
[41, 42, 43, 44]],
dtype=dtype, device=device)
mm = torch.sparse.mm if masked else torch.mm
blocksize = (2, 2) if layout in {torch.sparse_bsr, torch.sparse_bsc} else None
x = ref_x.to_sparse(layout=layout, blocksize=blocksize).requires_grad_(True)
y = ref_y.requires_grad_(True)
if layout is torch.sparse_bsr and not masked or layout is torch.sparse_bsc:
with self.assertRaisesRegex(
RuntimeError,
r"addmm: computation on (CPU|CUDA) is not implemented for Strided \+ Sparse(Bsr|Bsc) @ Strided"):
torch.autograd.gradcheck(mm, (x, y), fast_mode=fast_mode, masked=masked)
self.skipTest('NOT IMPL')
elif layout in {torch.sparse_csc, torch.sparse_bsr, torch.sparse_bsc} and masked:
with self.assertRaisesRegex(
RuntimeError,
r"(sparse_addmm_sparse_backward: unsupported combination of layouts,"
r" grad: Strided, mat1: Sparse(Csc|Bsr|Bsc), mat2: Strided"
r"|addmm: computation on (CPU|CUDA) is not implemented for "
r"Strided \+ Sparse(Csc|Bsr|Bsc) @ Strided without MKL)"):
torch.autograd.gradcheck(mm, (x, y), fast_mode=fast_mode, masked=masked)
self.skipTest('NOT IMPL')
else:
torch.autograd.gradcheck(mm, (x, y), fast_mode=fast_mode, masked=masked)
@onlyNativeDeviceTypes
@suppress_warnings
@ops(binary_ufuncs_with_sparse_support)
@all_sparse_layouts('layout', include_strided=False)
def test_binary_operation(self, layout, device, dtype, op):
if not op.supports_sparse_layout(layout):
self.skipTest(f'{layout} is not supported in `{op.name}` OpInfo definition. Skipping!')
for sample in op.sample_inputs_sparse(layout, device, dtype):
if validate_sample_input_sparse(op, sample, check_validate=False) is not sample:
# that is, the validation returns the sparse sample
# wrapped within ErrorInput instance
continue
t_inp, t_args, t_kwargs = sample.input, sample.args, sample.kwargs
batch_dim = t_inp.dim() - t_inp.dense_dim() - t_inp.sparse_dim()
result = op.op(t_inp, *t_args, **t_kwargs)
# Check rop(inp, ...).shape == inp.shape
self.assertEqual(result.shape, t_inp.shape)
# Check rop(inp, ...).sparse_dim() == inp.sparse_dim()
self.assertEqual(result.sparse_dim(), t_inp.sparse_dim())
# Check rop(inp, ...).dense_dim() == inp.dense_dim()
self.assertEqual(result.dense_dim(), t_inp.dense_dim())
# Check invariant rop(inp, ...).to_dense() == rop(inp.to_dense(), ...)
try:
dense = op.op(t_inp.to_dense(), *(t_args[0].to_dense(), *t_args[1:]), **t_kwargs)
except Exception as msg:
# this is strided op issue, so skipping the sample silently here
if "\"cpublas_axpy_impl\" not implemented for 'ComplexHalf'" in str(msg):
continue
raise
self.assertEqual(result, dense)
@onlyCPU
@all_sparse_layouts('layout', include_strided=True)
@dtypes(torch.double)
def test_to_sparse_identity(self, device, layout, dtype):
for dense_dim in range(4):
x_dense = torch.eye(dense_dim, dtype=dtype, device=device)
for sparse_dim_in in range(1, dense_dim):
x_sparse = x_dense.to_sparse(sparse_dim_in)
for sparse_dim_out in range(0, dense_dim):
if sparse_dim_out == sparse_dim_in:
self.assertTrue(x_sparse.to_sparse(sparse_dim_out).sparse_dim() == sparse_dim_out)
else:
with self.assertRaisesRegex(
RuntimeError,
r"to_sparse: conversion from Sparse to Sparse with sparse_dim argument !=self.sparse_dim\(\)"
" is not supported"):
x_sparse.to_sparse(sparse_dim_out)
@onlyNativeDeviceTypes
@suppress_warnings
@ops(like_fns_with_sparse_support)
@all_sparse_layouts('layout', include_strided=False)
def test_like_fns(self, layout, device, dtype, op):
for sample in op.sample_inputs_sparse(layout, device, dtype):
t_inp, t_args, t_kwargs = sample.input, sample.args, sample.kwargs
batch_dim = t_inp.dim() - t_inp.dense_dim() - t_inp.sparse_dim()
if t_inp.layout in {torch.sparse_bsr, torch.sparse_bsc}:
expected_blocksize = t_inp.values().shape[batch_dim + 1:batch_dim + 3]
else:
expected_blocksize = None
expected_dtype = t_kwargs.get('dtype', dtype)
expected_device = torch.device(t_kwargs.get('device', device))
expected_layout = t_kwargs.get('layout', layout)
result = op.op(t_inp, *t_args, **t_kwargs)
self.assertEqual(result.dtype, expected_dtype)
self.assertEqual(result.device.type, expected_device.type)
self.assertEqual(result.layout, expected_layout)
if result.layout in {torch.sparse_bsr, torch.sparse_bsc}:
result_batch_dim = result.dim() - result.dense_dim() - result.sparse_dim()
blocksize = result.values().shape[result_batch_dim + 1:result_batch_dim + 3]
self.assertEqual(blocksize, expected_blocksize)
# Check op(inp).shape == inp.shape
self.assertEqual(result.shape, t_inp.shape)
if expected_layout is torch.strided:
self.assertEqual(result.sparse_dim(), 0)
# Check op(inp, layout=torch.strided).dense_dim() == inp.dim()
self.assertEqual(result.dense_dim(), t_inp.dim())
elif expected_layout is torch.sparse_coo:
# Check op(inp, layout=torch.sparse_coo).sparse_dim() == batch_dim + inp.sparse_dim()
self.assertEqual(result.sparse_dim(), batch_dim + t_inp.sparse_dim())
# Check op(inp, layout=torch.sparse_coo).dense_dim() == inp.dense_dim()
self.assertEqual(result.dense_dim(), t_inp.dense_dim())
torch._validate_sparse_coo_tensor_args(result._indices(), result._values(), result.shape)
else:
# Check op(inp).sparse_dim() == inp.sparse_dim()
self.assertEqual(result.sparse_dim(), t_inp.sparse_dim())
# Check op(inp).dense_dim() == inp.dense_dim()
self.assertEqual(result.dense_dim(), t_inp.dense_dim())
if result.layout in {torch.sparse_csr, torch.sparse_bsr}:
compressed_indices, plain_indices = result.crow_indices(), result.col_indices()
else:
compressed_indices, plain_indices = result.ccol_indices(), result.row_indices()
torch._validate_sparse_compressed_tensor_args(compressed_indices, plain_indices, result.values(),
result.shape, result.layout)
@all_sparse_layouts('mask_layout', include_strided=False)
@onlyNativeDeviceTypes
@dtypes(*all_types_and_complex_and(torch.half, torch.bool, torch.bfloat16))
def test_sparse_mask(self, mask_layout, device, dtype):
input_layout = torch.strided
mask_dtype = torch.bool
for mask in self.generate_simple_inputs(mask_layout, dtype=mask_dtype, device=device,
enable_hybrid=False, enable_batch=False):
x = make_tensor(mask.shape, dtype=dtype, device=device).to_sparse(layout=input_layout)
result = x.sparse_mask(mask)
# Check invariant `x.sparse_mask(mask).<indices> == mask.<indices>`
if mask_layout is torch.sparse_coo:
self.assertEqual(result._indices(), mask._indices())
ones = torch.sparse_coo_tensor(mask._indices(),
torch.ones_like(mask._values(), dtype=x.dtype),
mask.shape,
is_coalesced=mask.is_coalesced())
elif mask_layout in {torch.sparse_csr, torch.sparse_bsr}:
self.assertEqual(result.crow_indices(), mask.crow_indices())
self.assertEqual(result.col_indices(), mask.col_indices())
ones = torch.sparse_compressed_tensor(mask.crow_indices(), mask.col_indices(),
torch.ones_like(mask.values(), dtype=x.dtype),
mask.shape, layout=mask.layout)
else:
self.assertEqual(result.ccol_indices(), mask.ccol_indices())
self.assertEqual(result.row_indices(), mask.row_indices())
ones = torch.sparse_compressed_tensor(mask.ccol_indices(), mask.row_indices(),
torch.ones_like(mask.values(), dtype=x.dtype),
mask.shape, layout=mask.layout)
# Check invariant:
# x.sparse_mask(mask).to_dense() == x.mul(sparse_xyz_tensor(<mask indices>,
# ones_like(<mask values>)).to_dense())
expected = x.mul(ones.to_dense())
self.assertEqual(result.to_dense(), expected)
# Check invariant `mask.to_dense().sparse_mask(mask) == mask`
result = mask.to_dense().sparse_mask(mask)
self.assertEqual(result, mask)
@all_sparse_layouts('layout', include_strided=False)
@parametrize("masked", [subtest(False, name='nonmasked'), subtest(True, name='masked')])
@parametrize("fast_mode", [subtest(False, name='slow'), subtest(True, name='fast')])
def test_as_sparse_gradcheck(self, layout, device, masked, fast_mode):
gradcheck = torch.sparse.as_sparse_gradcheck(torch.autograd.gradcheck)
sparse_compressed_layouts = {torch.sparse_csr, torch.sparse_csc, torch.sparse_bsr, torch.sparse_bsc}
def identity(x):
return x
for func in (torch.Tensor.to_dense,
torch.Tensor.sum,
identity,
torch.Tensor.to_sparse,
torch.Tensor.values,
):
for x in self.generate_simple_inputs(
layout,
device=device,
dtype=torch.float64,
# TODO: fix gh-104868 to enable batched samples:
enable_batch=layout not in sparse_compressed_layouts,
enable_hybrid=not (
layout in sparse_compressed_layouts and (
# FIXME: RuntimeError: sparse_mask(): the
# number of sparse dimensions in `self`
# should match that of the `mask`. Got
# `self.sparse_dim() == 3` !=
# `mask.sparse_dim() == 2
func.__name__ == 'sum'
# FIXME: RuntimeError: expected
# col_indices to be a contiguous tensor
# per batch
or func.__name__ == 'to_sparse'
))):
if layout is torch.sparse_coo and func.__name__ == 'values':
x = x.coalesce()
gradcheck(func, x.requires_grad_(True), masked=masked, fast_mode=fast_mode)
@onlyCPU
@all_sparse_layouts('layout', include_strided=False)
@dtypes(torch.double)
def test_dataloader(self, device, layout, dtype):
data = list(self.generate_simple_inputs(layout, device=device, dtype=dtype))
dataset = _SparseDataset(data)
loader = torch.utils.data.DataLoader(dataset, batch_size=None, num_workers=2)
loaded_data = list(loader)
self.assertEqual(data, loaded_data)
@onlyCPU
def test_invalid_blocksize(self):
# Blocksize should be a tuple/list/torch.Size containing two values
with self.assertRaisesRegex(RuntimeError, ".*blocksize.*, but got 1"):
torch.randn(1).to_sparse(blocksize=(1,))
with self.assertRaisesRegex(RuntimeError, ".*blocksize.*, but got 1"):
torch.randn(1).to_sparse(blocksize=[1])
with self.assertRaisesRegex(RuntimeError, ".*blocksize.*, but got 1"):
torch.randn(1).to_sparse(blocksize=torch.Size((1,)))
with self.assertRaisesRegex(RuntimeError, ".*blocksize.*, but got 3"):
torch.randn(1).to_sparse(blocksize=(1, 1, 1))
with self.assertRaisesRegex(RuntimeError, ".*blocksize.*, but got 3"):
torch.randn(1).to_sparse(blocksize=[1, 1, 1])
with self.assertRaisesRegex(RuntimeError, ".*blocksize.*, but got 3"):
torch.randn(1).to_sparse(blocksize=torch.Size((1, 1, 1)))
@unittest.skipIf(not torch.cuda.is_available(), 'requires cuda')
@onlyCPU
@all_sparse_layouts('layout', include_strided=True)
def test_constructor_pin_memory(self, device, layout):
"""Tests sparse_xyz_tensor(indices, values, pin_memory=True)
"""
self.assertEqual(device, "cpu")
for t in self.generate_simple_inputs(
layout, device=device, dtype=torch.float64,
enable_zero_sized=False, # pinning zero-sized tensors is a no-op
pin_memory=True,
enable_batch=False, # TODO: remove after gh-104868 is resolved
):
if layout is torch.sparse_coo:
self.assertTrue(t._indices().is_pinned())
self.assertTrue(t._values().is_pinned())
elif layout in {torch.sparse_csr, torch.sparse_bsr}:
self.assertTrue(t.crow_indices().is_pinned())
self.assertTrue(t.col_indices().is_pinned())
self.assertTrue(t.values().is_pinned())
elif layout in {torch.sparse_csc, torch.sparse_bsc}:
self.assertTrue(t.ccol_indices().is_pinned())
self.assertTrue(t.row_indices().is_pinned())
self.assertTrue(t.values().is_pinned())
elif layout is torch.strided:
pass
else:
assert 0 # unreachable
self.assertTrue(t.is_pinned())
@unittest.skipIf(not torch.cuda.is_available(), 'requires cuda')
@onlyCPU
@all_sparse_layouts('layout', include_strided=True)
def test_method_pin_memory(self, device, layout):
"""Tests sparse_xyz_tensor(indices, values, pin_memory=False).pin_memory()
"""
for t_ in self.generate_simple_inputs(
layout, device=device, dtype=torch.float64,
enable_zero_sized=False, # pinning zero-sized tensors is a no-op
pin_memory=False, # no pinning
enable_batch=False, # TODO: remove after gh-104868 is resolved
):
t = t_.pin_memory()
self.assertTrue(t.is_pinned())
# registering a non-pinned tensor with CUDA memory is a
# clone operation
self.assertFalse(t_.is_pinned())
# registering already pinned tensor with CUDA memory is an
# identity operation:
t2 = t.pin_memory()
self.assertTrue(t2 is t)
if layout is torch.sparse_coo:
self.assertTrue(t._indices().is_pinned())
self.assertTrue(t._values().is_pinned())
self.assertFalse(t_._indices().is_pinned())
self.assertFalse(t_._values().is_pinned())
elif layout in {torch.sparse_csr, torch.sparse_bsr}:
self.assertTrue(t.crow_indices().is_pinned())
self.assertTrue(t.col_indices().is_pinned())
self.assertTrue(t.values().is_pinned())
self.assertFalse(t_.crow_indices().is_pinned())
self.assertFalse(t_.col_indices().is_pinned())
self.assertFalse(t_.values().is_pinned())
elif layout in {torch.sparse_csc, torch.sparse_bsc}:
self.assertTrue(t.ccol_indices().is_pinned())
self.assertTrue(t.row_indices().is_pinned())
self.assertTrue(t.values().is_pinned())
self.assertFalse(t_.ccol_indices().is_pinned())
self.assertFalse(t_.row_indices().is_pinned())
self.assertFalse(t_.values().is_pinned())
elif layout is torch.strided:
pass
else:
assert 0 # unreachable
@unittest.skipIf(not torch.cuda.is_available(), 'requires cuda')
@onlyCPU
@all_sparse_layouts('layout', include_strided=True)
def test_constructor_pinned_memory(self, device, layout):
"""Tests sparse_xyz_tensor(indices.pin_memory(device), values.pin_memory(device))
"""
pin_memory_device = "cuda"
for t in self.generate_simple_inputs(
layout, device=device, dtype=torch.float64,
enable_zero_sized=False, # pinning zero-sized tensors is a no-op
pin_memory=None, # constructor does not specify pin_memory=...
members_pin_memory=True, # indices and values are pinned
enable_batch=False, # TODO: remove after gh-104868 is resolved
):
if layout is torch.sparse_coo:
self.assertTrue(t._indices().is_pinned())
self.assertTrue(t._values().is_pinned())
elif layout in {torch.sparse_csr, torch.sparse_bsr}:
self.assertTrue(t.crow_indices().is_pinned())
self.assertTrue(t.col_indices().is_pinned())
self.assertTrue(t.values().is_pinned())
elif layout in {torch.sparse_csc, torch.sparse_bsc}:
self.assertTrue(t.ccol_indices().is_pinned())
self.assertTrue(t.row_indices().is_pinned())
self.assertTrue(t.values().is_pinned())
elif layout is torch.strided:
pass
else:
assert 0 # unreachable
self.assertTrue(t.is_pinned())
@unittest.skipIf(not torch.cuda.is_available(), 'requires cuda')
@onlyCPU
@all_sparse_layouts('layout', include_strided=False)
def test_constructor_mismatched_pinned_memory(self, device, layout):
"""Test the failure to construct sparse tensor from indices and values
that have different pinning states.
"""
def generic_constructor(*args, **kwargs):
if layout in {torch.sparse_csr, torch.sparse_csc, torch.sparse_bsr, torch.sparse_bsc}:
kwargs.update(layout=layout)
return torch.sparse_compressed_tensor(*args, **kwargs)
elif layout is torch.sparse_coo:
return torch.sparse_coo_tensor(*args, **kwargs)
else:
raise NotImplementedError(layout)
for args, kwargs in self.generate_simple_inputs(
layout, device=device, dtype=torch.float64,
enable_zero_sized=False, # pinning zero-sized tensors is a no-op
enable_batch=False, # TODO: remove after gh-104868 is resolved
output_tensor=False):
# indices are pinned, values is a non-pinned tensor
args1 = (args[0].pin_memory(), *args[1:])
# indices are non-pinned, values is a pinned tensor
args2 = (*args[:-1], args[-1].pin_memory())
with self.assertRaisesRegex(
RuntimeError, r"memory pinning of \w*indices \(=1\) must match memory pinning of values \(=0\)"):
generic_constructor(*args1, **kwargs)
with self.assertRaisesRegex(
RuntimeError, r"memory pinning of \w*indices \(=0\) must match memory pinning of values \(=1\)"):
generic_constructor(*args2, **kwargs)
# e.g., TestSparseUnaryUfuncsCPU and TestSparseUnaryUfuncsCUDA
instantiate_device_type_tests(TestSparseUnaryUfuncs, globals(), except_for='meta')
instantiate_device_type_tests(TestSparseMaskedReductions, globals(), except_for='meta')
# e.g., TestSparseCPU and TestSparseCUDA
instantiate_device_type_tests(TestSparse, globals(), except_for='meta')
instantiate_device_type_tests(TestSparseAny, globals(), except_for='meta')
instantiate_parametrized_tests(TestSparseMeta)
instantiate_parametrized_tests(TestSparseLegacyAndDeprecation)
if __name__ == '__main__':
run_tests()