mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
711e5a6ceb
* Port THS to ATen.
The basic structure of the patch:
- All kernels in aten/src/THS got rewritten as native
functions in aten/src/ATen/native/sparse
I took the liberty to rename some of the kernels,
opting for a longer, more transparent names than
things like 'spaddcmul'.
- Instead of holding fields for sparse tensor in the TH
C struct THSTensor, they are now held in a C++ class
SparseTensorImpl (this explains why I had to do this
all in one go; I can't have *two* reps for sparse
tensors!)
Along the way, we change a key internal representation
invariant: an "empty" sparse tensor has dimI == 1 and
dimV == 0 (this is different from dimI == 0 and dimV == 0
we had before); this ensures that we maintain the invariant
that dim == dimI + dimV. "Scalar" sparse tensors are
made illegal, because there really is no way to properly
express them in COO format.
- Because we haven't ported THCS or any of the traditional
dense TH implementations, there is a new set of adapter
functions in native/LegacyBridge.cpp exclusively devoted
to deciding whether or not to go to the new native implementation
or back to the legacy TH binding (prefixed with th_).
The intent is that when everything gets ported, we can
delete this file.
- I've kept the stubs for all the THS functions, but they now all
error if you try to actually call them. Eventually, we should
replace these with calls to ATen so that everything keeps
working.
- I gobbled up SparseMM (SparseMM.cpp is no more). It was tasty.
There are some miscellaneous improvements which were needed for other
changes in this patch:
- There is now AT_FORALL_SCALAR_TYPES_EXCEPT_HALF, which does what
it says on the tin.
- axpy templated function moved to TH/BlasUtils.h, there's a new macro
which lets you easily forward to all of the TH functions. We also expose
THBlas_copy. I'm not terribly pleased with these functions but
they seem to serve a purpose they need.
- New method on Tensor to get TensorImpl*, unsafeGetTensorImpl
- accessor() is now this-const, since const-correctness on Tensor is a lie
- New toSparse()/toDense() methods on Type; now you can call these
directly without having to manually apply at::toSparse/toDense
on the Backend and then running toBackend yourself.
Changes to the kernels:
- Previously, the whole body of all kernels was compiled for
every supported scalar type. In our new implementation,
the scalar dispatch has been pushed into the smallest extent
which (1) is not in a type loop and (2) requires statically
knowing the scalar type. These sites all use
AT_DISPATCH_ALL_TYPES. I tried to use lambdas as much as
possible, but sometimes it was not possible when a OpenMP
pragma was used.
- Anywhere we tested if the nDimension of a tensor was zero,
we replaced with a test that numel is zero. Because, as we
known, nDimension of zero-size tensors in TH is zero, and
that's wrong wrong wrong (and not done this way in ATen).
Some subtleties:
- Places where previously fastget1d was used, I now use a
TensorAccessor. However, you have to be careful about grabbing
the accessor, because sometimes you will be accessor'ing
indices/values and they are empty, which means they will
be *1D* ("oh, aren't indices always 2D?" Nope. Nyet.)
So, essentially, it is only safe to grab an accessor *after*
you have checked that nnz != 0. All of these shenanigans
will go away when we properly support zero-size dimensions.
A few places, we test for this case just by wrapping the loop
in a conditional on nnz. Some other places this is not so easy,
so we instead short-circuit the function with a special case for
when nnz == 0 (usually, these implementations are degenerate).
- There is a very subtle but important difference between
_sparse_get_impl(self)->indices() and self._indices();
the latter may return a view! This is because nnz is
not guaranteed to match the dimensions of indices/values;
you can "truncate" a sparse tensor by setting the nnz.
Actually, I think this is not a good idea and we should
enforce a stronger invariant, but for this patch I slavishly
adhere to the old ways, and as such I have to be very
careful if I want to resize something, I had better use
the former and not the latter.
- I had to reimplement broadcasting by hand (thus the s_
and non-s_ functions in the sparse native files). There
is a very important distinction between foo_out and foo_,
so it is important that the LegacyBridge function always
call to the lower layer, and not try to avoid boilerplate
by calling to another LegacyBridge function first.
I did NOT put broadcasting in LegacyBridge (even though,
ultimately, that's where it must live), because the th_
functions which are invoked from LegacyBridge handle
broadcasting themselves, and I don't want to broadcast
twice.
- Sparse function MUST explicitly specify the Type they
dispatch from, otherwise Variable wrapping/unwrapping will
not work correctly. If you use _get_sparse_impl, that is
sufficient to levy this requirement.
- The "has native" tests in LegacyBridge.cpp are not 100%,
because some of the functions are mixed dense-sparse functions,
and so you can't just say, "Oh, if it's sparse and CPU, call
the native sparse implementation." This is handled on a
case by case basis. There is some especially complex
logic for add(), which has dense-dense, sparse-sparse
and dense-sparse implementations.
- I added some uses of SparseTensorRef in native_functions.yaml,
but you will notice that these are all on native_* functions,
and not the actual, top-level functions. So the SparseTensorRef
is purely documentary (helping you not call the wrong overload)
but there is no magic; we do the wrapping ourselves the hard
way. (This is in constrast to the TH binding code which is magical.)
Except for _sparse_mask; _sparse_mask is magical.
- There is a raw_copy_sparse_ method, which is really my way of
getting around the fact that copy_ has never been implemented
for sparse tensors (even before this patch), but there IS a
super secret, internal way of doing these copies that the THS
code used, and which I needed to get my hands on when I did this
port. We should refactor so that either (a) copy_ does support
sparse-sparse copy natively, or (b) we do this other ways.
- Irritatingly, I must explicitly resize_as_ before copy_ into
a tensor. This was not the case with THTensor_(copy) but I don't
have any direct binding that doesn't have this requirement.
- For some reason, the sparse tensor constructor accepts a scalar
tensor for the values tensor. This is kind of weird because
you always need an nnz-dimension. However, the old code supported
this and just expanded it into a 1D size 0 tensor; so we need some
explicit code to do this.
There are maybe a bit more AT_ASSERTs in some of the kernels
than is wise. I added them all when I was debugging and was
loathe to remove them.
Some last mile fixes after this commit went into PR
- Move expand outside of dispatch so autograd works (it used to be inside and then we lost all of the recorded broadcasts).
- Hack to duplicate the derivatives for our now two definitions TH and native. Mercifully the derivatives are short.
- Apparently, TH has a special case to make foo_ functions method only, and if you don't do this the Python arg parsing is wrong. We carefully work around this in the native bindings
- Apply DCE to a test_jit case, fixes wobbling due to DCE trick in tracing
- Update test_function's output
- Some last mile fixes for dispatch confusion in sparse_coo_tensor functions.
- New simplified regression test based on failures I saw in ONNX
- Increase tolerance on super resolution test
- More robust dynamic_type normalization, fixes ONNX bug.
The dynamic_type situation is very delicate; probably need
to stop having both Scalar and real.
- Make new_with_tensor_sparse more CUDA safe
- Note about CUDA-safety in SparseTensorImpl
- Rename dimI/dimV to sparseDims/denseDims.
- Make localScalar on SparseTensorImpl work.
- Make numel uniformly supported on all types, not just dense
types
- Add tests for is_nonzero() method (which exercises localScalar)
- Disable constant JIT autogenerated tests, which are fragile and broken
by this change, but being fixed in a parallel track.
Signed-off-by: Edward Z. Yang <ezyang@fb.com>
1040 lines
39 KiB
Python
1040 lines
39 KiB
Python
import torch
|
|
from torch import sparse
|
|
|
|
import itertools
|
|
import random
|
|
import unittest
|
|
from common import TestCase, run_tests
|
|
from common_cuda import TEST_CUDA
|
|
from test_torch import TestTorch
|
|
from numbers import Number
|
|
|
|
|
|
def cpu_only(inner):
|
|
def outer(self, *args, **kwargs):
|
|
if self.is_cuda:
|
|
raise unittest.SkipTest("Test is CPU-only")
|
|
inner(self, *args, **kwargs)
|
|
return outer
|
|
|
|
|
|
def cuda_only(inner):
|
|
def outer(self, *args, **kwargs):
|
|
if not self.is_cuda:
|
|
raise unittest.SkipTest("Test is GPU-only")
|
|
inner(self, *args, **kwargs)
|
|
return outer
|
|
|
|
|
|
class TestSparse(TestCase):
|
|
|
|
def setUp(self):
|
|
# These parameters control the various ways we can run the test.
|
|
# We will subclass and override this method to implement CUDA
|
|
# tests
|
|
self.is_cuda = False
|
|
self.is_uncoalesced = False
|
|
self.IndexTensor = torch.LongTensor
|
|
self.ValueTensor = torch.DoubleTensor
|
|
self.SparseTensor = torch.sparse.DoubleTensor
|
|
super(TestSparse, self).setUp()
|
|
|
|
def _gen_sparse(self, d, nnz, with_size):
|
|
# TODO: Consider implementing this in the CUDA case by directly
|
|
# performing the operations on the GPU. You won't be able to
|
|
# use torch.rand/torch.randn in this case because they are
|
|
# CPU-only. If you do this, you can remove the is_cuda branch
|
|
# at the end.
|
|
#
|
|
# If you do this, be sure to update assert_uncoalesced too
|
|
|
|
if isinstance(with_size, Number):
|
|
with_size = [with_size] * d
|
|
|
|
if self.is_uncoalesced:
|
|
# We want to generate a tensor with a lot of uncoalesced
|
|
# entries to stress test whether or not we handle this
|
|
# (subtle) case correctly
|
|
v_size = [nnz * 2] + list(with_size[d:])
|
|
v = torch.randn(*v_size)
|
|
r = torch.rand(d, nnz)
|
|
# Repeat the indexes, so every position shows up twice
|
|
i = torch.cat([r, r], dim=1) * \
|
|
torch.Tensor(with_size[:d]).repeat(nnz * 2, 1).transpose(0, 1)
|
|
i = i.type(torch.LongTensor)
|
|
x = torch.sparse.DoubleTensor(i, v, torch.Size(with_size))
|
|
self.assert_uncoalesced(x)
|
|
else:
|
|
# Generate a sparse tensor with d sparse dimensions; the
|
|
# rest the dimensions with_size[d:] are dense.
|
|
v_size = [nnz] + list(with_size[d:])
|
|
v = torch.randn(*v_size)
|
|
i = torch.rand(d, nnz) * \
|
|
torch.Tensor(with_size[:d]).repeat(nnz, 1).transpose(0, 1)
|
|
i = i.type(torch.LongTensor)
|
|
x = torch.sparse.DoubleTensor(i, v, torch.Size(with_size))
|
|
|
|
if self.is_cuda:
|
|
return x.cuda(), i.cuda(), v.cuda()
|
|
else:
|
|
return x, i.clone(), v.clone()
|
|
|
|
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()
|
|
# Strategy: construct a new sparse tensor with the raw value
|
|
# field overwritten to a tensor of ones, coalesce it, and then
|
|
# check if any value entries are > 1 (which indicates that the
|
|
# original was uncoalesced.)
|
|
i = x._indices().clone()
|
|
v = x._values().clone().fill_(1)
|
|
y = torch.sparse.DoubleTensor(i, v, x.size())
|
|
z = self.safeCoalesce(y)
|
|
assert (z._values() > 1).sum() > 0
|
|
|
|
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 self.ValueTensor(*args, **kwargs).normal_()
|
|
|
|
def test_basic(self):
|
|
x, i, v = self._gen_sparse(3, 10, 100)
|
|
|
|
self.assertEqual(i, x._indices())
|
|
self.assertEqual(v, x._values())
|
|
|
|
x, i, v = self._gen_sparse(3, 10, [100, 100, 100])
|
|
self.assertEqual(i, x._indices())
|
|
self.assertEqual(v, x._values())
|
|
self.assertEqual(x.ndimension(), 3)
|
|
self.assertEqual(self.safeCoalesce(x)._nnz(), 10)
|
|
for i in range(3):
|
|
self.assertEqual(x.size(i), 100)
|
|
|
|
# Make sure that coalesce handles duplicate indices correctly
|
|
i = self.IndexTensor([[9, 0, 0, 0, 8, 1, 1, 1, 2, 7, 2, 2, 3, 4, 6, 9]])
|
|
v = self.ValueTensor([[idx**2, idx] for idx in range(i.size(1))])
|
|
x = self.SparseTensor(i, v, torch.Size([10, 2]))
|
|
self.assertEqual(self.safeCoalesce(x)._nnz(), 9)
|
|
|
|
# Make sure we can access empty indices / values
|
|
x = self.SparseTensor()
|
|
self.assertEqual(x._indices().numel(), 0)
|
|
self.assertEqual(x._values().numel(), 0)
|
|
|
|
def test_ctor_size_checks(self):
|
|
indices = self.IndexTensor([
|
|
[0, 0, 0],
|
|
[0, 3, 0],
|
|
[0, 0, 0],
|
|
[0, 0, 0],
|
|
])
|
|
values = self.ValueTensor([2, 1, 3, 4])
|
|
|
|
# indices inconsistent with size
|
|
self.assertRaises(
|
|
RuntimeError,
|
|
lambda: self.SparseTensor(indices, values, torch.Size([2, 1, 1])))
|
|
|
|
# values inconsistent with size
|
|
values = self.ValueTensor([
|
|
[2, 1, 2, 1],
|
|
[1, 0, 5, 2],
|
|
])
|
|
self.assertRaises(
|
|
RuntimeError,
|
|
lambda: self.SparseTensor(indices, values, torch.Size([2, 4, 2, 1])))
|
|
|
|
def test_to_dense(self):
|
|
i = self.IndexTensor([
|
|
[0, 1, 2, 2],
|
|
[0, 0, 0, 3],
|
|
[0, 0, 1, 4],
|
|
])
|
|
v = self.ValueTensor([2, 1, 3, 4])
|
|
x = self.SparseTensor(i, v, torch.Size([3, 4, 5]))
|
|
res = self.ValueTensor([
|
|
[[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]],
|
|
])
|
|
|
|
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 test_shared(self):
|
|
i = self.IndexTensor([[2]])
|
|
v = self.ValueTensor([5])
|
|
x = self.SparseTensor(i, v, torch.Size([3]))
|
|
v[0] = 6
|
|
self.assertEqual(self.ValueTensor([0, 0, 6]), self.safeToDense(x))
|
|
i[0][0] = 0
|
|
self.assertEqual(self.ValueTensor([6, 0, 0]), self.safeToDense(x))
|
|
|
|
def test_to_dense_hybrid(self):
|
|
i = self.IndexTensor([
|
|
[0, 1, 2, 2],
|
|
[0, 0, 0, 3],
|
|
])
|
|
v = self.ValueTensor([[2, 3], [1, 2], [3, 4], [4, 5]])
|
|
x = self.SparseTensor(i, v, torch.Size([3, 4, 2]))
|
|
res = self.ValueTensor([
|
|
[[2, 3],
|
|
[0, 0],
|
|
[0, 0],
|
|
[0, 0]],
|
|
[[1, 2],
|
|
[0, 0],
|
|
[0, 0],
|
|
[0, 0]],
|
|
[[3, 4],
|
|
[0, 0],
|
|
[0, 0],
|
|
[4, 5]],
|
|
])
|
|
|
|
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 test_contig(self):
|
|
i = self.IndexTensor([
|
|
[1, 0, 35, 14, 39, 6, 71, 66, 40, 27],
|
|
[92, 31, 62, 50, 22, 65, 89, 74, 56, 34],
|
|
])
|
|
v = self.ValueTensor([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
|
|
x = self.SparseTensor(i, v, torch.Size([100, 100]))
|
|
exp_i = self.IndexTensor([
|
|
[0, 1, 6, 14, 27, 35, 39, 40, 66, 71],
|
|
[31, 92, 65, 50, 34, 62, 22, 56, 74, 89],
|
|
])
|
|
exp_v = self.ValueTensor([2, 1, 6, 4, 10, 3, 5, 9, 8, 7])
|
|
x = self.safeCoalesce(x)
|
|
self.assertEqual(exp_i, x._indices())
|
|
self.assertEqual(exp_v, x._values())
|
|
|
|
i = self.IndexTensor([
|
|
[2, 0, 2, 1],
|
|
[0, 0, 3, 0],
|
|
[1, 0, 4, 0],
|
|
])
|
|
v = self.ValueTensor([3, 2, 4, 1])
|
|
x = self.SparseTensor(i, v, torch.Size([3, 4, 5]))
|
|
exp_i = self.IndexTensor([
|
|
[0, 1, 2, 2],
|
|
[0, 0, 0, 3],
|
|
[0, 0, 1, 4],
|
|
])
|
|
exp_v = self.ValueTensor([2, 1, 3, 4])
|
|
|
|
x = self.safeCoalesce(x)
|
|
self.assertEqual(exp_i, x._indices())
|
|
self.assertEqual(exp_v, x._values())
|
|
|
|
# Duplicate indices
|
|
i = self.IndexTensor([
|
|
[0, 0, 2, 0],
|
|
[0, 0, 3, 0],
|
|
[0, 0, 4, 0],
|
|
])
|
|
v = self.ValueTensor([3, 2, 4, 1])
|
|
x = self.SparseTensor(i, v, torch.Size([3, 4, 5]))
|
|
exp_i = self.IndexTensor([
|
|
[0, 2],
|
|
[0, 3],
|
|
[0, 4],
|
|
])
|
|
exp_v = self.ValueTensor([6, 4])
|
|
|
|
x = self.safeCoalesce(x)
|
|
self.assertEqual(exp_i, x._indices())
|
|
self.assertEqual(exp_v, x._values())
|
|
|
|
def test_contig_hybrid(self):
|
|
i = self.IndexTensor([
|
|
[1, 0, 35, 14, 39, 6, 71, 66, 40, 27],
|
|
[92, 31, 62, 50, 22, 65, 89, 74, 56, 34],
|
|
])
|
|
v = self.ValueTensor([
|
|
[1, 2], [2, 3], [3, 4], [4, 5], [5, 6],
|
|
[6, 7], [7, 8], [8, 9], [9, 10], [10, 11],
|
|
])
|
|
x = self.SparseTensor(i, v, torch.Size([100, 100, 2]))
|
|
exp_i = self.IndexTensor([
|
|
[0, 1, 6, 14, 27, 35, 39, 40, 66, 71],
|
|
[31, 92, 65, 50, 34, 62, 22, 56, 74, 89],
|
|
])
|
|
exp_v = self.ValueTensor([
|
|
[2, 3], [1, 2], [6, 7], [4, 5], [10, 11],
|
|
[3, 4], [5, 6], [9, 10], [8, 9], [7, 8],
|
|
])
|
|
x = self.safeCoalesce(x)
|
|
self.assertEqual(exp_i, x._indices())
|
|
self.assertEqual(exp_v, x._values())
|
|
|
|
i = self.IndexTensor([
|
|
[2, 0, 2, 1],
|
|
[0, 0, 3, 0],
|
|
[1, 0, 4, 0],
|
|
])
|
|
v = self.ValueTensor([[3, 3, 3], [2, 2, 2], [4, 4, 4], [1, 1, 1]])
|
|
x = self.SparseTensor(i, v, torch.Size([3, 4, 5, 3]))
|
|
exp_i = self.IndexTensor([
|
|
[0, 1, 2, 2],
|
|
[0, 0, 0, 3],
|
|
[0, 0, 1, 4],
|
|
])
|
|
exp_v = self.ValueTensor([[2, 2, 2], [1, 1, 1], [3, 3, 3], [4, 4, 4]])
|
|
|
|
x = self.safeCoalesce(x)
|
|
self.assertEqual(exp_i, x._indices())
|
|
self.assertEqual(exp_v, x._values())
|
|
|
|
# Duplicate indices
|
|
i = self.IndexTensor([
|
|
[0, 0, 2, 0],
|
|
[0, 0, 3, 0],
|
|
[0, 0, 4, 0],
|
|
])
|
|
v = self.ValueTensor([[3, 2, 3], [2, 1, 1], [4, 3, 4], [1, 1, 1]])
|
|
x = self.SparseTensor(i, v, torch.Size([3, 4, 5, 3]))
|
|
exp_i = self.IndexTensor([
|
|
[0, 2],
|
|
[0, 3],
|
|
[0, 4],
|
|
])
|
|
exp_v = self.ValueTensor([[6, 4, 5], [4, 3, 4]])
|
|
|
|
x = self.safeCoalesce(x)
|
|
self.assertEqual(exp_i, x._indices())
|
|
self.assertEqual(exp_v, x._values())
|
|
|
|
def test_clone(self):
|
|
x, _, _ = self._gen_sparse(4, 20, 5)
|
|
if self.is_uncoalesced:
|
|
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())
|
|
|
|
@cuda_only
|
|
def test_cuda_empty(self):
|
|
x = torch.sparse.FloatTensor(2, 3, 4)
|
|
y = x.cuda(0)
|
|
self.assertEqual(x._sparseDims(), y._sparseDims())
|
|
self.assertEqual(x._denseDims(), y._denseDims())
|
|
x = y.cpu()
|
|
self.assertEqual(y._sparseDims(), x._sparseDims())
|
|
self.assertEqual(y._denseDims(), x._denseDims())
|
|
|
|
def test_transpose(self):
|
|
x = self._gen_sparse(4, 20, 5)[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)
|
|
|
|
def test_transpose_coalesce_invariant(self):
|
|
# If a sparse tensor is coalesced, its transpose should be the same
|
|
# If a sparse tensor is uncoalesed, its transpose should be the same
|
|
x_coalesced = self._gen_sparse(2, 3, 4)[0].coalesce()
|
|
x_indices = x_coalesced._indices()
|
|
x_values = x_coalesced._values()
|
|
|
|
y_uncoalesced = self.SparseTensor(
|
|
torch.cat([x_indices, x_indices], dim=1),
|
|
torch.cat([x_values, x_values]),
|
|
x_coalesced.size())
|
|
|
|
self.assertTrue(x_coalesced.is_coalesced())
|
|
self.assertFalse(y_uncoalesced.is_coalesced())
|
|
|
|
self.assertTrue(x_coalesced.transpose(0, 1).is_coalesced())
|
|
self.assertFalse(y_uncoalesced.transpose(0, 1).is_coalesced())
|
|
|
|
x_coalesced.transpose_(0, 1)
|
|
y_uncoalesced.transpose_(0, 1)
|
|
self.assertTrue(x_coalesced.is_coalesced())
|
|
self.assertFalse(y_uncoalesced.is_coalesced())
|
|
|
|
def test_t_empty(self):
|
|
x = self.SparseTensor(2, 3)
|
|
x.t_()
|
|
self.assertEqual(torch.Size([3, 2]), x.size())
|
|
self.assertEqual(0, x._indices().numel())
|
|
self.assertEqual(0, x._values().numel())
|
|
self.assertEqual(x._sparseDims(), 2)
|
|
self.assertEqual(x._denseDims(), 0)
|
|
|
|
x = self.SparseTensor(2, 3)
|
|
y = x.t()
|
|
self.assertEqual(torch.Size([3, 2]), y.size())
|
|
self.assertEqual(0, y._indices().numel())
|
|
self.assertEqual(0, y._values().numel())
|
|
self.assertEqual(x._sparseDims(), 2)
|
|
self.assertEqual(x._denseDims(), 0)
|
|
|
|
def test_add_zeros(self):
|
|
def test_shape(sparse_dims, sizes):
|
|
x, _, _ = self._gen_sparse(sparse_dims, 20, sizes)
|
|
zeros = torch.zeros(sizes, layout=torch.sparse_coo).to(x.device)
|
|
r1 = zeros + x
|
|
r2 = x + zeros
|
|
self.assertEqual(r1, x)
|
|
self.assertEqual(r2, x)
|
|
|
|
test_shape(1, [1])
|
|
test_shape(4, [3, 17, 19, 5])
|
|
test_shape(2, [3, 17, 19, 5])
|
|
|
|
@cpu_only
|
|
def test_mm(self):
|
|
def test_shape(di, dj, dk):
|
|
x, _, _ = self._gen_sparse(2, 20, [di, dj])
|
|
t = torch.randn(di, dk)
|
|
y = torch.randn(dj, dk)
|
|
alpha = random.random()
|
|
beta = random.random()
|
|
|
|
res = torch.addmm(alpha, t, beta, x, y)
|
|
expected = torch.addmm(alpha, t, beta, self.safeToDense(x), y)
|
|
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)
|
|
test_shape(100, 1000, 200)
|
|
test_shape(64, 10000, 300)
|
|
|
|
@cpu_only
|
|
def test_saddmm(self):
|
|
def test_shape(di, dj, dk):
|
|
x = self._gen_sparse(2, 20, [di, dj])[0]
|
|
t = self._gen_sparse(2, 20, [di, dk])[0]
|
|
y = torch.randn(dj, dk)
|
|
alpha = random.random()
|
|
beta = random.random()
|
|
|
|
res = torch.saddmm(alpha, t, beta, x, y)
|
|
expected = torch.addmm(alpha, self.safeToDense(t), beta, self.safeToDense(x), y)
|
|
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)
|
|
test_shape(1000, 100, 100)
|
|
test_shape(3000, 64, 300)
|
|
|
|
def test_dsmm(self):
|
|
def test_shape(di, dj, dk):
|
|
x = self._gen_sparse(2, 20, [di, dj])[0]
|
|
y = self.randn(dj, dk)
|
|
|
|
res = torch.dsmm(x, y)
|
|
expected = torch.mm(self.safeToDense(x), y)
|
|
self.assertEqual(res, expected)
|
|
|
|
test_shape(7, 5, 3)
|
|
test_shape(1000, 100, 100)
|
|
test_shape(3000, 64, 300)
|
|
|
|
def test_hsmm(self):
|
|
def test_shape(di, dj, dk):
|
|
x = self._gen_sparse(2, 20, [di, dj])[0]
|
|
y = self.randn(dj, dk)
|
|
|
|
res = torch.hsmm(x, y)
|
|
# TODO: use self.safeToDense(), but this triggers
|
|
# https://github.com/pytorch/pytorch/issues/3170
|
|
expected = torch.mm(x.to_dense(), y)
|
|
self.assertEqual(res.to_dense(), expected)
|
|
|
|
test_shape(7, 5, 3)
|
|
test_shape(1000, 100, 100)
|
|
test_shape(3000, 64, 300)
|
|
|
|
def _test_spadd_shape(self, shape_i, shape_v=None):
|
|
shape = shape_i + (shape_v or [])
|
|
x, _, _ = self._gen_sparse(len(shape_i), 10, shape)
|
|
y = self.randn(*shape)
|
|
r = random.random()
|
|
|
|
res = torch.add(y, r, x)
|
|
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)
|
|
y.transpose_(0, len(s) - 1)
|
|
r = random.random()
|
|
|
|
res = torch.add(y, r, x)
|
|
expected = y + r * self.safeToDense(x)
|
|
|
|
self.assertEqual(res, expected)
|
|
|
|
def test_spadd(self):
|
|
self._test_spadd_shape([5, 6])
|
|
self._test_spadd_shape([10, 10, 10])
|
|
self._test_spadd_shape([50, 30, 20])
|
|
self._test_spadd_shape([5, 5, 5, 5, 5, 5])
|
|
|
|
def test_spadd_hybrid(self):
|
|
self._test_spadd_shape([5, 6], [2, 3])
|
|
self._test_spadd_shape([10, 10, 10], [3])
|
|
self._test_spadd_shape([50, 30, 20], [2])
|
|
self._test_spadd_shape([5, 5, 5, 5, 5, 5], [2])
|
|
|
|
def test_norm(self):
|
|
x, _, _ = self._gen_sparse(3, 10, 100)
|
|
y = x.coalesce()
|
|
self.assertEqual(x.norm(), y._values().norm())
|
|
|
|
def _test_basic_ops_shape(self, shape_i, shape_v=None):
|
|
shape = shape_i + (shape_v or [])
|
|
x1, _, _ = self._gen_sparse(len(shape_i), 9, shape)
|
|
x2, _, _ = self._gen_sparse(len(shape_i), 12, shape)
|
|
|
|
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)
|
|
|
|
# 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())
|
|
self.assertEqual(self.safeToDense(y), expected)
|
|
|
|
self.assertFalse(x1.is_coalesced())
|
|
y = x1.coalesce()
|
|
z = x1.coalesce()
|
|
self.assertFalse(x1.is_coalesced())
|
|
self.assertTrue(y.is_coalesced())
|
|
self.assertEqual(x1, y)
|
|
# check that coalesce is out of place
|
|
y._values().add_(1)
|
|
self.assertEqual(z._values() + 1, y._values())
|
|
|
|
def test_basic_ops(self):
|
|
self._test_basic_ops_shape([5, 6])
|
|
self._test_basic_ops_shape([10, 10, 10])
|
|
self._test_basic_ops_shape([50, 30, 20])
|
|
self._test_basic_ops_shape([5, 5, 5, 5, 5, 5])
|
|
|
|
def test_basic_ops_hybrid(self):
|
|
self._test_basic_ops_shape([5, 6], [2, 3])
|
|
self._test_basic_ops_shape([10, 10, 10], [3])
|
|
self._test_basic_ops_shape([50, 30, 20], [2])
|
|
self._test_basic_ops_shape([5, 5, 5, 5, 5, 5], [2])
|
|
|
|
def _test_sparse_mask_shape(self, shape_i, shape_v=None):
|
|
shape = shape_i + (shape_v or [])
|
|
x1, _, _ = self._gen_sparse(len(shape_i), 9, shape)
|
|
x2, _, _ = self._gen_sparse(len(shape_i), 12, shape)
|
|
|
|
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)
|
|
|
|
def _test_sparse_mask_fixed(self):
|
|
i = self.IndexTensor([
|
|
[1, 3, 0, 4],
|
|
[2, 1, 2, 3],
|
|
])
|
|
v = self.ValueTensor([1, 2, 3, 4])
|
|
x = self.SparseTensor(i, v, torch.Size([5, 4])).coalesce()
|
|
dense = self.ValueTensor([
|
|
[1, 2, 3, 4],
|
|
[5, 6, 7, 8],
|
|
[9, 10, 11, 12],
|
|
[13, 14, 15, 16],
|
|
[17, 18, 19, 20],
|
|
])
|
|
exp_v = self.ValueTensor([7, 14, 3, 20])
|
|
res = dense._sparse_mask(x)
|
|
expected = self.SparseTensor(i, exp_v, torch.Size([5, 4]))
|
|
self.assertEqual(res, expected)
|
|
|
|
def test_sparse_mask(self):
|
|
self._test_sparse_mask_fixed()
|
|
|
|
self._test_sparse_mask_shape([5, 6])
|
|
self._test_sparse_mask_shape([10, 10, 10])
|
|
self._test_sparse_mask_shape([50, 30, 20])
|
|
self._test_sparse_mask_shape([5, 5, 5, 5, 5, 5])
|
|
|
|
def _test_zeros(self, shape, out_shape_i, out_shape_v=None):
|
|
out_shape = out_shape_i + (out_shape_v or [])
|
|
for nnz in [9, 12]:
|
|
out, _, _ = self._gen_sparse(len(out_shape_i), nnz, out_shape)
|
|
torch.zeros(*shape, out=out)
|
|
self.assertEqual(tuple(out.size()), tuple(shape))
|
|
self.assertTrue(out._indices().numel() == out._values().numel() == 0)
|
|
self.assertEqual(out._nnz(), 0)
|
|
self.assertEqual(out._sparseDims(), len(shape))
|
|
self.assertEqual(out._denseDims(), 0)
|
|
|
|
def test_zeros(self):
|
|
i_shapes = [2, 3, 4]
|
|
v_shapes = [3, 4, 5, 6]
|
|
for i_dim in range(1, len(i_shapes) + 1):
|
|
for v_dim in range(len(v_shapes) + 1):
|
|
self._test_zeros([2, 3, 4], i_shapes[:i_dim], v_shapes[:v_dim])
|
|
|
|
def _test_zeros_like(self, 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 [9, 12]:
|
|
t, _, _ = self._gen_sparse(len(template_shape_i), nnz, template_shape)
|
|
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._sparseDims(), len(template_shape_i))
|
|
self.assertEqual(res._denseDims(), len(template_shape_v))
|
|
|
|
def test_zeros_like(self):
|
|
i_shapes = [2, 3, 4]
|
|
v_shapes = [3, 4, 5, 6]
|
|
for i_dim in range(1, len(i_shapes) + 1):
|
|
for v_dim in range(len(v_shapes) + 1):
|
|
self._test_zeros_like(i_shapes[:i_dim], v_shapes[:v_dim])
|
|
|
|
def _test_sparse_mask_hybrid_fixed(self):
|
|
i = self.IndexTensor([
|
|
[1, 3, 0, 4],
|
|
[2, 1, 2, 3],
|
|
])
|
|
v = self.ValueTensor([[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.SparseTensor(i, v, torch.Size([5, 4, 2])).coalesce()
|
|
dense = self.ValueTensor([
|
|
[[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._sparse_mask(x)
|
|
exp_v = self.ValueTensor([[7, 9], [14, 1], [3, 3], [20, 1]])
|
|
expected = self.SparseTensor(i, exp_v, torch.Size([5, 4, 2]))
|
|
self.assertEqual(res, expected)
|
|
|
|
def test_sparse_variable_methods(self):
|
|
# TODO: delete when tensor/variable are merged
|
|
from torch.autograd import Variable
|
|
i = self.IndexTensor([[0, 1, 1], [2, 0, 2]])
|
|
v = self.ValueTensor([3, 4, 5])
|
|
sparse_mat = self.SparseTensor(i, v, torch.Size([2, 3]))
|
|
sparse_var = Variable(sparse_mat)
|
|
|
|
to_test_one_arg = {
|
|
'zeros_like': lambda x: torch.zeros_like(x),
|
|
'transpose': lambda x: x.transpose(0, 1),
|
|
'transpose_': lambda x: x.transpose_(0, 1),
|
|
't': lambda x: x.t(),
|
|
't_': lambda x: x.t_(),
|
|
'div': lambda x: x.div(2),
|
|
'div_': lambda x: x.div_(2),
|
|
'pow': lambda x: x.pow(2),
|
|
'_nnz': lambda x: x._nnz(),
|
|
'is_coalesced': lambda x: x.is_coalesced(),
|
|
'coalesce': lambda x: x.coalesce(),
|
|
'to_dense': lambda x: x.to_dense(),
|
|
'_sparseDims': lambda x: x._sparseDims(),
|
|
'_denseDims': lambda x: x._denseDims(),
|
|
'norm': lambda x: x.norm(),
|
|
}
|
|
|
|
for test_name, test_fn in to_test_one_arg.items():
|
|
var1 = sparse_var.clone()
|
|
tensor1 = sparse_mat.clone()
|
|
|
|
out_var = test_fn(var1)
|
|
out_tensor = test_fn(tensor1)
|
|
|
|
if isinstance(out_tensor, int) or isinstance(out_tensor, bool):
|
|
if not isinstance(out_var, int) and not isinstance(out_var, bool):
|
|
check_var = out_var.data[0]
|
|
else:
|
|
check_var = out_var
|
|
self.assertEqual(out_var, out_tensor)
|
|
continue
|
|
|
|
# Assume output is variable / tensor
|
|
self.assertEqual(test_fn(var1).data, test_fn(tensor1),
|
|
test_name)
|
|
|
|
i = self.IndexTensor([[0, 0, 1], [1, 2, 1]])
|
|
v = self.ValueTensor([3, 3, 4])
|
|
sparse_mat2 = self.SparseTensor(i, v, torch.Size([2, 3]))
|
|
sparse_var2 = Variable(sparse_mat2)
|
|
|
|
to_test_two_arg = {
|
|
'sub': lambda x, y: x.sub(y),
|
|
'sub_': lambda x, y: x.sub_(y),
|
|
'mul': lambda x, y: x.mul(y),
|
|
'mul_': lambda x, y: x.mul_(y),
|
|
}
|
|
|
|
for test_name, test_fn in to_test_two_arg.items():
|
|
var1 = sparse_var.clone()
|
|
var2 = sparse_var2.clone()
|
|
tensor1 = sparse_mat.clone()
|
|
tensor2 = sparse_mat2.clone()
|
|
self.assertEqual(test_fn(var1, var2).data,
|
|
test_fn(tensor1, tensor2), test_name)
|
|
|
|
to_test_mixed = [
|
|
# test name, lambda expression, should_run_when_cuda
|
|
('sspaddmm', lambda sp, de: sp.sspaddmm(sp, de), False),
|
|
('sspaddmm_b', lambda sp, de: sp.sspaddmm(2, sp, de), False),
|
|
('sspaddmm_b_a', lambda sp, de: sp.sspaddmm(3, 2, sp, de), False),
|
|
('addmm', lambda sp, de: de.addmm(sp, de), True),
|
|
# TODO: This looks like a typo
|
|
('addmm_', lambda sp, de: de.addmm(sp, de), True),
|
|
('mm', lambda sp, de: torch.mm(sp, de), True),
|
|
('mm_out', lambda sp, de: torch.mm(sp, de, out=de), True),
|
|
]
|
|
|
|
i = self.IndexTensor([[0, 0, 1, 2, 2], [1, 2, 1, 0, 1]])
|
|
v = self.ValueTensor([3, 3, 4, 1, 2])
|
|
sparse_mat = self.SparseTensor(i, v, torch.Size([3, 3]))
|
|
sparse_var = Variable(sparse_mat)
|
|
dense_mat = sparse_mat.to_dense().random_(0, 5)
|
|
dense_var = Variable(dense_mat)
|
|
|
|
for test_name, test_fn, test_cuda in to_test_mixed:
|
|
if sparse_var.is_cuda and not test_cuda:
|
|
continue
|
|
sp_var = sparse_var.clone()
|
|
de_var = dense_var.clone()
|
|
sp_mat = sparse_mat.clone()
|
|
de_mat = dense_mat.clone()
|
|
self.assertEqual(test_fn(sp_var, de_var).data,
|
|
test_fn(sp_mat, de_mat), test_name)
|
|
|
|
def test_sparse_mask_hybrid(self):
|
|
self._test_sparse_mask_hybrid_fixed()
|
|
|
|
self._test_sparse_mask_shape([5, 6], [2, 3])
|
|
self._test_sparse_mask_shape([10, 10, 10], [3])
|
|
self._test_sparse_mask_shape([50, 30, 20], [2])
|
|
self._test_sparse_mask_shape([5, 5, 5, 5, 5, 5], [2])
|
|
|
|
def test_sparse_add_coalesce(self):
|
|
i = self.IndexTensor([[1, 2, 1]])
|
|
v = self.ValueTensor([3, 4, 5])
|
|
x = self.SparseTensor(i, v, torch.Size([3]))
|
|
y = self.SparseTensor(i, v, torch.Size([3]))
|
|
z = x + y
|
|
|
|
self.assertFalse(z._indices().numel() != 2 and z.is_coalesced())
|
|
|
|
@cuda_only
|
|
def test_storage_not_null(self):
|
|
x = torch.cuda.sparse.FloatTensor(2)
|
|
self.assertNotEqual(x.get_device(), -1)
|
|
|
|
@cuda_only
|
|
@unittest.skipIf(torch.cuda.device_count() < 2, "only one GPU detected")
|
|
def test_same_gpu(self):
|
|
i = self.IndexTensor([[2]]).cuda(1)
|
|
v = self.ValueTensor([5]).cuda(1)
|
|
x = self.SparseTensor(i, v, torch.Size([3]), device=1)
|
|
self.assertEqual(x.get_device(), 1)
|
|
self.assertEqual(x._values().get_device(), 1)
|
|
self.assertEqual(x._indices().get_device(), 1)
|
|
|
|
x = self.SparseTensor(3, device=1)
|
|
self.assertEqual(x.get_device(), 1)
|
|
self.assertEqual(x._values().get_device(), 1)
|
|
self.assertEqual(x._indices().get_device(), 1)
|
|
|
|
v = self.ValueTensor([5]).cuda(0)
|
|
self.assertRaises(RuntimeError, lambda: self.SparseTensor(i, v, torch.Size([3])))
|
|
|
|
def _test_new_device(self, size, device):
|
|
with torch.cuda.device(device):
|
|
x = torch.cuda.sparse.DoubleTensor(*size)
|
|
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)
|
|
|
|
@cuda_only
|
|
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)
|
|
|
|
@cuda_only
|
|
@unittest.skipIf(torch.cuda.device_count() < 2, "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)
|
|
|
|
def test_new(self):
|
|
x, indices, values = self._gen_sparse(3, 10, 100)
|
|
if not x.is_cuda:
|
|
# CUDA sparse tensors currently requires the size to be
|
|
# specified if nDimV > 0
|
|
self.assertEqual(x.new(indices, values), x)
|
|
self.assertEqual(x.new(indices, values, x.size()), x)
|
|
|
|
@cpu_only # not really, but we only really want to run this once
|
|
def test_factory(self):
|
|
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
|
|
dtype = torch.float64
|
|
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])
|
|
values = torch.tensor([1.], dtype=dtype) if use_tensor_val else 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)
|
|
|
|
def test_factory_size_check(self):
|
|
indices = self.IndexTensor([[1, 2], [0, 2]])
|
|
values = self.ValueTensor([.5, .5])
|
|
sizes = torch.Size([2, 3])
|
|
with self.assertRaisesRegex(RuntimeError, "sizes is inconsistent with indices"):
|
|
self.SparseTensor(indices, values, sizes)
|
|
|
|
indices = self.IndexTensor([[1, 2], [0, 2]])
|
|
values = self.ValueTensor([[1, 1, 1], [1, 1, 1]])
|
|
sizes = torch.Size([3, 3, 2])
|
|
with self.assertRaisesRegex(RuntimeError, "values and sizes are inconsistent"):
|
|
self.SparseTensor(indices, values, sizes)
|
|
|
|
def test_factory_empty_indices(self):
|
|
device = 'cuda' if self.is_cuda else 'cpu'
|
|
tensor = torch.sparse_coo_tensor([], [], torch.Size([]), device=device)
|
|
expected_indices = torch.tensor([], dtype=torch.long, device=device)
|
|
self.assertEqual(tensor._indices(), expected_indices)
|
|
|
|
@cpu_only
|
|
def test_factory_type_inference(self):
|
|
t = torch.sparse_coo_tensor(torch.tensor(([0], [2])), torch.tensor([1.], dtype=torch.float32))
|
|
self.assertEqual(torch.float32, t.dtype)
|
|
t = torch.sparse_coo_tensor(torch.tensor(([0], [2])), torch.tensor([1.], dtype=torch.float64))
|
|
self.assertEqual(torch.float64, t.dtype)
|
|
t = torch.sparse_coo_tensor(torch.tensor(([0], [2])), torch.tensor([1]))
|
|
self.assertEqual(torch.int64, t.dtype)
|
|
|
|
@cuda_only
|
|
def test_factory_device_type_inference(self):
|
|
# both indices/values are CUDA
|
|
shape = (1, 3)
|
|
for indices_device in ['cuda', 'cpu']:
|
|
for values_device in ['cuda', 'cpu']:
|
|
for sparse_device in ['cuda', 'cpu', None]:
|
|
t = torch.sparse_coo_tensor(torch.tensor(([0], [2]), device=indices_device),
|
|
torch.tensor([1.], device=values_device),
|
|
(1, 3), device=sparse_device)
|
|
should_be_cuda = sparse_device == 'cuda' or (sparse_device is None and values_device == 'cuda')
|
|
self.assertEqual(should_be_cuda, t.is_cuda)
|
|
|
|
@cpu_only
|
|
def test_factory_copy(self):
|
|
# both correct
|
|
indices = torch.tensor(([0], [2]), dtype=torch.int64)
|
|
values = torch.tensor([1.], dtype=torch.float64)
|
|
sparse_tensor = torch.sparse_coo_tensor(indices, values, dtype=torch.float64)
|
|
self.assertEqual(indices.data_ptr(), sparse_tensor._indices().data_ptr())
|
|
self.assertEqual(values.data_ptr(), sparse_tensor._values().data_ptr())
|
|
|
|
# only indices correct
|
|
indices = torch.tensor(([0], [2]), dtype=torch.int64)
|
|
values = torch.tensor([1.], dtype=torch.float32)
|
|
sparse_tensor = torch.sparse_coo_tensor(indices, values, dtype=torch.float64)
|
|
self.assertEqual(indices.data_ptr(), sparse_tensor._indices().data_ptr())
|
|
self.assertNotEqual(values.data_ptr(), sparse_tensor._values().data_ptr())
|
|
|
|
# only values correct
|
|
indices = torch.tensor(([0], [2]), dtype=torch.int32)
|
|
values = torch.tensor([1.], dtype=torch.float64)
|
|
sparse_tensor = torch.sparse_coo_tensor(indices, values, dtype=torch.float64)
|
|
self.assertNotEqual(indices.data_ptr(), sparse_tensor._indices().data_ptr())
|
|
self.assertEqual(values.data_ptr(), sparse_tensor._values().data_ptr())
|
|
|
|
# neither correct
|
|
indices = torch.tensor(([0], [2]), dtype=torch.int32)
|
|
values = torch.tensor([1.], dtype=torch.float32)
|
|
sparse_tensor = torch.sparse_coo_tensor(indices, values, dtype=torch.float64)
|
|
self.assertNotEqual(indices.data_ptr(), sparse_tensor._indices().data_ptr())
|
|
self.assertNotEqual(values.data_ptr(), sparse_tensor._values().data_ptr())
|
|
|
|
@cpu_only # not really, but we only really want to run this once
|
|
def test_dtypes(self):
|
|
all_sparse_dtypes = [dtype for dtype in torch.testing.get_all_dtypes() if dtype != torch.float16]
|
|
TestTorch._test_dtypes(self, all_sparse_dtypes, torch.sparse_coo, torch.device('cpu'))
|
|
if torch.cuda.is_available():
|
|
TestTorch._test_dtypes(self, all_sparse_dtypes, torch.sparse_coo, torch.device('cuda:0'))
|
|
|
|
@cpu_only # not really, but we only really want to run this once
|
|
def test_empty_full(self):
|
|
all_sparse_dtypes = [dtype for dtype in torch.testing.get_all_dtypes() if dtype != torch.float16]
|
|
TestTorch._test_empty_full(self, all_sparse_dtypes, torch.sparse_coo, torch.device('cpu'))
|
|
if torch.cuda.device_count() > 0:
|
|
TestTorch._test_empty_full(self, all_sparse_dtypes, torch.sparse_coo, -1)
|
|
TestTorch._test_empty_full(self, all_sparse_dtypes, torch.sparse_coo, torch.device('cuda:0'))
|
|
|
|
def test_is_sparse(self):
|
|
x = torch.randn(3, 3)
|
|
self.assertFalse(x.is_sparse)
|
|
|
|
x = self.SparseTensor()
|
|
self.assertTrue(x.is_sparse)
|
|
|
|
def test_resize_as(self):
|
|
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
|
|
# _sparseDims and _denseDims match.
|
|
self.assertEqual(t, t + y)
|
|
|
|
do_test(self.SparseTensor())
|
|
|
|
def test_is_nonzero(self):
|
|
self.assertTrue(torch.sparse_coo_tensor(([0],), 1., (1,)).is_nonzero())
|
|
self.assertFalse(torch.sparse_coo_tensor(([0],), 0., (1,)).is_nonzero())
|
|
self.assertFalse(torch.sparse_coo_tensor(([0], [0]), 0., (1, 1)).is_nonzero())
|
|
self.assertFalse(torch.sparse_coo_tensor(([0, 0],), (0., 0.), (1,)).is_nonzero())
|
|
self.assertFalse(torch.sparse_coo_tensor(([0, 0],), (-1., 1.), (1,)).is_nonzero())
|
|
# NB: We should test "scalar" sparse tensors, but they don't actually
|
|
# work at the moment (in principle, they should)
|
|
|
|
|
|
class TestUncoalescedSparse(TestSparse):
|
|
def setUp(self):
|
|
super(TestUncoalescedSparse, self).setUp()
|
|
self.is_uncoalesced = True
|
|
|
|
|
|
@unittest.skipIf(not TEST_CUDA, 'CUDA not available')
|
|
class TestCudaSparse(TestSparse):
|
|
def setUp(self):
|
|
super(TestCudaSparse, self).setUp()
|
|
self.is_cuda = True
|
|
self.IndexTensor = torch.cuda.LongTensor
|
|
self.ValueTensor = torch.cuda.DoubleTensor
|
|
self.SparseTensor = torch.cuda.sparse.DoubleTensor
|
|
|
|
|
|
@unittest.skipIf(not TEST_CUDA, 'CUDA not available')
|
|
class TestCudaUncoalescedSparse(TestCudaSparse):
|
|
def setUp(self):
|
|
super(TestCudaUncoalescedSparse, self).setUp()
|
|
self.is_uncoalesced = True
|
|
|
|
if __name__ == '__main__':
|
|
run_tests()
|