Caffe2 - Python API
A deep learning, cross platform ML framework
test_sparse.py
1 import torch
2 from torch import sparse
3 
4 import itertools
5 import functools
6 import random
7 import unittest
8 from common_utils import TestCase, run_tests, skipIfRocm, do_test_dtypes, do_test_empty_full, load_tests
9 from common_cuda import TEST_CUDA
10 from numbers import Number
11 from torch.autograd.gradcheck import gradcheck
12 
13 # load_tests from common_utils is used to automatically filter tests for
14 # sharding on sandcastle. This line silences flake warnings
15 load_tests = load_tests
16 
17 
18 def cpu_only(inner):
19  @functools.wraps(inner)
20  def outer(self, *args, **kwargs):
21  if self.is_cuda:
22  raise unittest.SkipTest("Test is CPU-only")
23  inner(self, *args, **kwargs)
24  return outer
25 
26 
27 def cuda_only(inner):
28  @functools.wraps(inner)
29  def outer(self, *args, **kwargs):
30  if not self.is_cuda:
31  raise unittest.SkipTest("Test is GPU-only")
32  inner(self, *args, **kwargs)
33  return outer
34 
35 
37 
38  def setUp(self):
39  # These parameters control the various ways we can run the test.
40  # We will subclass and override this method to implement CUDA
41  # tests
42  self.is_cuda = False
43  self.is_uncoalesced = False
44  self.device = 'cpu'
45  self.value_dtype = torch.float64
46  self.index_tensor = lambda *args: torch.tensor(*args, dtype=torch.int64, device=self.device)
47  self.value_empty = lambda *args: torch.empty(*args, dtype=self.value_dtype, device=self.device)
48  self.value_tensor = lambda *args: torch.tensor(*args, dtype=self.value_dtype, device=self.device)
49 
50  def sparse_empty_factory(*args, **kwargs):
51  kwargs['dtype'] = kwargs.get('dtype', self.value_dtype)
52  kwargs['layout'] = kwargs.get('laytout', torch.sparse_coo)
53  kwargs['device'] = kwargs.get('device', self.device)
54  return torch.empty(*args, **kwargs)
55  self.sparse_empty = sparse_empty_factory
56 
57  def sparse_tensor_factory(*args, **kwargs):
58  kwargs['dtype'] = kwargs.get('dtype', self.value_dtype)
59  kwargs['device'] = kwargs.get('device', self.device)
60  return torch.sparse_coo_tensor(*args, **kwargs)
61  self.sparse_tensor = sparse_tensor_factory
62  self.legacy_sparse_tensor = torch.sparse.DoubleTensor
63  super(TestSparse, self).setUp()
64 
65  def _gen_sparse(self, sparse_dim, nnz, with_size):
66  if isinstance(with_size, Number):
67  with_size = [with_size] * sparse_dim
68 
69  x, i, v = self.genSparseTensor(with_size, sparse_dim, nnz, self.is_uncoalesced, self.device)
70 
71  if self.is_uncoalesced:
72  self.assert_uncoalesced(x)
73 
74  return x, i, v
75 
76  def assert_uncoalesced(self, x):
77  """
78  Test if a CPU tensor is uncoalesced. This is used to ensure
79  correctness of the uncoalesced tensor generation algorithm.
80  """
81  assert not x.is_coalesced()
82  existing_indices = set()
83  for i in range(x._nnz()):
84  index = str(x._indices()[:, i])
85  if index in existing_indices:
86  return True
87  else:
88  existing_indices.add(index)
89 
90  def randn(self, *args, **kwargs):
91  """
92  Variant of torch.randn that also works in the TEST_CUDA case.
93  """
94  # TODO: Put this in torch.cuda.randn
95  return self.value_empty(*args, **kwargs).normal_()
96 
97  @skipIfRocm # ROCm stack doesn't like the x + x call
98  def test_print(self):
99  shape_sparse_dim_nnz = [
100  ((), 0, 2),
101  ((0,), 0, 10),
102  ((2,), 0, 3),
103  ((100, 3), 1, 3),
104  ((100, 20, 3), 2, 0),
105  ((10, 0, 3), 0, 3),
106  ((10, 0, 3), 0, 0),
107  ]
108 
109  printed = []
110  for shape, sparse_dim, nnz in shape_sparse_dim_nnz:
111  indices_shape = torch.Size((sparse_dim, nnz))
112  values_shape = torch.Size((nnz,) + shape[sparse_dim:])
113  printed.append("# shape: {}".format(torch.Size(shape)))
114  printed.append("# nnz: {}".format(nnz))
115  printed.append("# sparse_dim: {}".format(sparse_dim))
116  printed.append("# indices shape: {}".format(indices_shape))
117  printed.append("# values shape: {}".format(values_shape))
118 
119  indices = torch.arange(indices_shape.numel(), dtype=self.index_tensor(0).dtype,
120  device=self.device).view(indices_shape)
121  for d in range(sparse_dim):
122  indices[d].clamp_(max=(shape[d] - 1)) # make it valid index
123  if self.is_uncoalesced and indices.numel() > 0:
124  indices[:, -1] = indices[:, 0] # make it uncoalesced
125  values_numel = values_shape.numel()
126  values = torch.arange(values_numel, dtype=self.value_dtype,
127  device=self.device).view(values_shape).div_(values_numel / 2.)
128  sp_tensor = self.sparse_tensor(indices, values, shape)
129 
130  dtypes = [torch.int32]
131  if values.dtype == torch.double:
132  dtypes.append(torch.float)
133  else:
134  dtypes.append(torch.double)
135  for dtype in dtypes:
136  printed.append("########## {} ##########".format(dtype))
137  x = sp_tensor.detach().to(dtype)
138  printed.append("# sparse tensor")
139  printed.append(str(x))
140  if x.dtype.is_floating_point:
141  printed.append("# after requires_grad_")
142  printed.append(str(x.requires_grad_()))
143  printed.append("# after addition")
144  printed.append(str(x + x))
145  printed.append("# _indices")
146  printed.append(str(x._indices()))
147  printed.append("# _values")
148  printed.append(str(x._values()))
149  printed.append('')
150  self.assertExpected('\n'.join(printed))
151 
152  def test_basic(self):
153  def test_shape(sparse_dims, nnz, with_size):
154  if isinstance(with_size, Number):
155  with_size = [with_size] * sparse_dims
156  x, i, v = self._gen_sparse(sparse_dims, nnz, with_size)
157  self.assertEqual(i, x._indices())
158  self.assertEqual(v, x._values())
159  self.assertEqual(x.ndimension(), len(with_size))
160  self.assertEqual(self.safeCoalesce(x)._nnz(), nnz)
161  self.assertEqual(list(x.size()), with_size)
162 
163  # Test .indices() and .values()
164  if self.is_uncoalesced:
165  with self.assertRaisesRegex(RuntimeError, "Cannot get indices on an uncoalesced tensor"):
166  x.indices()
167  with self.assertRaisesRegex(RuntimeError, "Cannot get values on an uncoalesced tensor"):
168  x.values()
169  else:
170  self.assertEqual(x.indices(), x._indices())
171  self.assertEqual(x.values(), x._values())
172 
173  test_shape(3, 10, 100)
174  test_shape(3, 10, [100, 100, 100])
175  test_shape(3, 10, [100, 100, 100, 5, 5, 5, 0])
176  test_shape(3, 0, [0, 0, 100, 5, 5, 5, 0])
177 
178  # Make sure that coalesce handles duplicate indices correctly
179  i = self.index_tensor([[9, 0, 0, 0, 8, 1, 1, 1, 2, 7, 2, 2, 3, 4, 6, 9]])
180  v = self.value_tensor([[idx**2, idx] for idx in range(i.size(1))])
181  x = self.sparse_tensor(i, v, torch.Size([10, 2]))
182  self.assertEqual(self.safeCoalesce(x)._nnz(), 9)
183 
184  # Make sure we can access empty indices / values
185  x = self.legacy_sparse_tensor()
186  self.assertEqual(x._indices().numel(), 0)
187  self.assertEqual(x._values().numel(), 0)
188 
189  def test_coalecce(self):
190  for empty_i, empty_v, empty_nnz in itertools.product([True, False], repeat=3):
191  sparse_size = [] if empty_i else [2, 1]
192  dense_size = [1, 0, 2] if empty_v else [1, 2]
193  nnz = 0 if empty_nnz else 5
194 
195  t, _, _ = self._gen_sparse(len(sparse_size), nnz, sparse_size + dense_size)
196  self.safeCoalesce(t) # this tests correctness
197 
198  def test_ctor_size_checks(self):
199  indices = self.index_tensor([
200  [0, 0, 0],
201  [0, 3, 0],
202  [0, 0, 0],
203  [0, 0, 0],
204  ])
205  values = self.value_tensor([2, 1, 3, 4])
206 
207  # indices inconsistent with size
208  self.assertRaises(
209  RuntimeError,
210  lambda: self.sparse_tensor(indices, values, torch.Size([2, 1, 1])))
211 
212  # values inconsistent with size
213  values = self.value_tensor([
214  [2, 1, 2, 1],
215  [1, 0, 5, 2],
216  ])
217  self.assertRaises(
218  RuntimeError,
219  lambda: self.sparse_tensor(indices, values, torch.Size([2, 4, 2, 1])))
220 
221  def test_to_dense(self):
222  def test_tensor(x, res):
223  x.to_dense() # Tests triple to_dense for memory corruption
224  x.to_dense()
225  x.to_dense()
226  self.assertEqual(res, x.to_dense())
227  self.assertEqual(res, self.safeToDense(x))
228 
229  def fn(x):
230  return x.to_dense()
231  x.requires_grad_(True)
232  gradcheck(fn, (x,), check_sparse_nnz=True)
233 
234  i = self.index_tensor([
235  [0, 1, 2, 2],
236  [0, 0, 0, 3],
237  [0, 0, 1, 4],
238  ])
239  v = self.value_tensor([2, 1, 3, 4])
240  x = self.sparse_tensor(i, v, torch.Size([3, 4, 5]))
241  res = self.value_tensor([
242  [[2, 0, 0, 0, 0],
243  [0, 0, 0, 0, 0],
244  [0, 0, 0, 0, 0],
245  [0, 0, 0, 0, 0]],
246  [[1, 0, 0, 0, 0],
247  [0, 0, 0, 0, 0],
248  [0, 0, 0, 0, 0],
249  [0, 0, 0, 0, 0]],
250  [[0, 3, 0, 0, 0],
251  [0, 0, 0, 0, 0],
252  [0, 0, 0, 0, 0],
253  [0, 0, 0, 0, 4]],
254  ])
255  test_tensor(x, res)
256 
257  i = self.index_tensor([
258  [0, 1, 2, 2],
259  [0, 0, 0, 3],
260  [0, 0, 1, 4],
261  ])
262  v = self.value_empty(4, 0)
263  x = self.sparse_tensor(i, v, torch.Size([3, 4, 5, 0]))
264  res = self.value_empty(3, 4, 5, 0)
265  test_tensor(x, res)
266 
267  def test_to_sparse(self):
268  shape = [10, 5, 19, 8]
269  max_nnz = 1
270  for dim, dim_sz in enumerate(shape, 1):
271  max_nnz *= dim_sz
272  rnnz = torch.randint(2, max_nnz, (1,)).item()
273  for nnz in [0, 1, rnnz]:
274  expected, _, _ = self._gen_sparse(dim, nnz, shape)
275  d = expected.to_dense()
276  result = d.to_sparse(dim)
277  self.assertEqual(d, result.to_dense()) # == not implemented for sparse tensors yet
278  self.assertEqual(expected.size(), result.size())
279  self.assertEqual(dim, result.sparse_dim())
280 
281  sp, _, _ = self._gen_sparse(2, 10, [3, 3, 3])
282  self.assertRaises(RuntimeError, lambda: sp.to_sparse())
283 
284  def test_scalar(self):
285  # tensor with value
286  a = self.sparse_tensor(self.index_tensor([]).unsqueeze(1), 12.3, [])
287  self.assertEqual(1, a._values().numel())
288  self.assertEqual(a, a.clone())
289  a_coalesced = a.coalesce()
290  self.assertTrue(a_coalesced.is_coalesced())
291  self.assertEqual(self.value_tensor(12.3), a.to_dense())
292  self.assertEqual(a, a.to_dense().to_sparse())
293 
294  # tensor with multiple values
295  a = self.sparse_tensor(self.index_tensor([]).unsqueeze(1).expand(0, 2), [12.3, 12.3], [])
296  self.assertEqual(2, a._values().numel())
297  self.assertEqual(a, a.clone())
298  a_coalesced = a.coalesce()
299  self.assertTrue(a_coalesced.is_coalesced())
300  self.assertEqual(self.value_tensor(12.3 * 2), a.to_dense())
301  self.assertEqual(a, a.to_dense().to_sparse())
302 
303  # tensor without value
304  a = self.sparse_empty(())
305  self.assertEqual(0, a._values().numel())
306  self.assertEqual(a, a.clone())
307  a_coalesced = a.coalesce()
308  self.assertTrue(a_coalesced.is_coalesced())
309  self.assertEqual(self.value_tensor(0), a.to_dense())
310  self.assertEqual(a, a.to_dense().to_sparse())
311 
312  def test_shared(self):
313  i = self.index_tensor([[2]])
314  v = self.value_tensor([5])
315  x = self.sparse_tensor(i, v, torch.Size([3]))
316  v[0] = 6
317  self.assertEqual(self.value_tensor([0, 0, 6]), self.safeToDense(x))
318  i[0][0] = 0
319  self.assertEqual(self.value_tensor([6, 0, 0]), self.safeToDense(x))
320 
321  i = self.index_tensor([[2]])
322  v = self.value_empty(1, 0)
323  x = self.sparse_tensor(i, v, torch.Size([3, 0]))
324  i[0][0] = 0
325  self.assertEqual(self.value_empty(3, 0), self.safeToDense(x))
326 
327  def test_to_dense_hybrid(self):
328  def test_tensor(x, res):
329  x.to_dense() # Tests double to_dense for memory corruption
330  x.to_dense()
331  x.to_dense()
332  self.assertEqual(res, x.to_dense())
333  self.assertEqual(res, self.safeToDense(x))
334 
335  def fn(x):
336  return x.to_dense()
337  x.requires_grad_(True)
338  gradcheck(fn, (x,), check_sparse_nnz=True)
339 
340  i = self.index_tensor([
341  [0, 1, 2, 2],
342  [0, 0, 0, 3],
343  ])
344  v = self.value_tensor([[2, 3], [1, 2], [3, 4], [4, 5]])
345  x = self.sparse_tensor(i, v, torch.Size([3, 4, 2]))
346  res = self.value_tensor([
347  [[2, 3],
348  [0, 0],
349  [0, 0],
350  [0, 0]],
351  [[1, 2],
352  [0, 0],
353  [0, 0],
354  [0, 0]],
355  [[3, 4],
356  [0, 0],
357  [0, 0],
358  [4, 5]],
359  ])
360  test_tensor(x, res)
361 
362  i = self.index_tensor([
363  [0, 1, 2, 2],
364  [0, 0, 0, 3],
365  ])
366  v = self.value_empty(4, 2, 0)
367  x = self.sparse_tensor(i, v, torch.Size([3, 4, 2, 0]))
368  res = self.value_empty(3, 4, 2, 0)
369  test_tensor(x, res)
370 
371  def test_contig(self):
372  def test_tensor(x, exp_i, exp_v):
373  x = self.safeCoalesce(x)
374  self.assertEqual(exp_i, x._indices())
375  self.assertEqual(exp_v, x._values())
376 
377  i = self.index_tensor([
378  [1, 0, 35, 14, 39, 6, 71, 66, 40, 27],
379  [92, 31, 62, 50, 22, 65, 89, 74, 56, 34],
380  ])
381  v = self.value_tensor([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
382  x = self.sparse_tensor(i, v, torch.Size([100, 100]))
383  exp_i = self.index_tensor([
384  [0, 1, 6, 14, 27, 35, 39, 40, 66, 71],
385  [31, 92, 65, 50, 34, 62, 22, 56, 74, 89],
386  ])
387  exp_v = self.value_tensor([2, 1, 6, 4, 10, 3, 5, 9, 8, 7])
388  test_tensor(x, exp_i, exp_v)
389 
390  i = self.index_tensor([
391  [2, 0, 2, 1],
392  [0, 0, 3, 0],
393  [1, 0, 4, 0],
394  ])
395  v = self.value_tensor([3, 2, 4, 1])
396  x = self.sparse_tensor(i, v, torch.Size([3, 4, 5]))
397  exp_i = self.index_tensor([
398  [0, 1, 2, 2],
399  [0, 0, 0, 3],
400  [0, 0, 1, 4],
401  ])
402  exp_v = self.value_tensor([2, 1, 3, 4])
403  test_tensor(x, exp_i, exp_v)
404 
405  i = self.index_tensor([
406  [2, 0, 2, 1],
407  [0, 0, 3, 0],
408  [1, 0, 4, 0],
409  ])
410  v = self.value_empty(4, 0)
411  x = self.sparse_tensor(i, v, torch.Size([3, 4, 5, 0]))
412  exp_i = self.index_tensor([
413  [0, 1, 2, 2],
414  [0, 0, 0, 3],
415  [0, 0, 1, 4],
416  ])
417  exp_v = self.value_empty(4, 0)
418  test_tensor(x, exp_i, exp_v)
419 
420  # Duplicate indices
421  i = self.index_tensor([
422  [0, 0, 2, 0],
423  [0, 0, 3, 0],
424  [0, 0, 4, 0],
425  ])
426  v = self.value_tensor([3, 2, 4, 1])
427  x = self.sparse_tensor(i, v, torch.Size([3, 4, 5]))
428  exp_i = self.index_tensor([
429  [0, 2],
430  [0, 3],
431  [0, 4],
432  ])
433  exp_v = self.value_tensor([6, 4])
434  test_tensor(x, exp_i, exp_v)
435 
436  i = self.index_tensor([
437  [0, 0, 2, 0],
438  [0, 0, 3, 0],
439  [0, 0, 4, 0],
440  ])
441  v = self.value_empty(4, 0)
442  x = self.sparse_tensor(i, v, torch.Size([3, 4, 5, 0]))
443  exp_i = self.index_tensor([
444  [0, 2],
445  [0, 3],
446  [0, 4],
447  ])
448  exp_v = self.value_empty(2, 0)
449  test_tensor(x, exp_i, exp_v)
450 
451  def test_contig_hybrid(self):
452  def test_tensor(x, exp_i, exp_v):
453  x = self.safeCoalesce(x)
454  self.assertEqual(exp_i, x._indices())
455  self.assertEqual(exp_v, x._values())
456 
457  i = self.index_tensor([
458  [1, 0, 35, 14, 39, 6, 71, 66, 40, 27],
459  [92, 31, 62, 50, 22, 65, 89, 74, 56, 34],
460  ])
461  v = self.value_tensor([
462  [1, 2], [2, 3], [3, 4], [4, 5], [5, 6],
463  [6, 7], [7, 8], [8, 9], [9, 10], [10, 11],
464  ])
465  x = self.sparse_tensor(i, v, torch.Size([100, 100, 2]))
466  exp_i = self.index_tensor([
467  [0, 1, 6, 14, 27, 35, 39, 40, 66, 71],
468  [31, 92, 65, 50, 34, 62, 22, 56, 74, 89],
469  ])
470  exp_v = self.value_tensor([
471  [2, 3], [1, 2], [6, 7], [4, 5], [10, 11],
472  [3, 4], [5, 6], [9, 10], [8, 9], [7, 8],
473  ])
474  test_tensor(x, exp_i, exp_v)
475 
476  i = self.index_tensor([
477  [2, 0, 2, 1],
478  [0, 0, 3, 0],
479  [1, 0, 4, 0],
480  ])
481  v = self.value_tensor([[3, 3, 3], [2, 2, 2], [4, 4, 4], [1, 1, 1]])
482  x = self.sparse_tensor(i, v, torch.Size([3, 4, 5, 3]))
483  exp_i = self.index_tensor([
484  [0, 1, 2, 2],
485  [0, 0, 0, 3],
486  [0, 0, 1, 4],
487  ])
488  exp_v = self.value_tensor([[2, 2, 2], [1, 1, 1], [3, 3, 3], [4, 4, 4]])
489  test_tensor(x, exp_i, exp_v)
490 
491  i = self.index_tensor([
492  [2, 0, 2, 1],
493  [0, 0, 3, 0],
494  [1, 0, 4, 0],
495  ])
496  v = self.value_empty(4, 3, 0)
497  x = self.sparse_tensor(i, v, torch.Size([3, 4, 5, 3, 0]))
498  exp_i = self.index_tensor([
499  [0, 1, 2, 2],
500  [0, 0, 0, 3],
501  [0, 0, 1, 4],
502  ])
503  exp_v = self.value_empty(4, 3, 0)
504  test_tensor(x, exp_i, exp_v)
505 
506  # Duplicate indices
507  i = self.index_tensor([
508  [0, 0, 2, 0],
509  [0, 0, 3, 0],
510  [0, 0, 4, 0],
511  ])
512  v = self.value_tensor([[3, 2, 3], [2, 1, 1], [4, 3, 4], [1, 1, 1]])
513  x = self.sparse_tensor(i, v, torch.Size([3, 4, 5, 3]))
514  exp_i = self.index_tensor([
515  [0, 2],
516  [0, 3],
517  [0, 4],
518  ])
519  exp_v = self.value_tensor([[6, 4, 5], [4, 3, 4]])
520  test_tensor(x, exp_i, exp_v)
521 
522  i = self.index_tensor([
523  [0, 0, 2, 0],
524  [0, 0, 3, 0],
525  [0, 0, 4, 0],
526  ])
527  v = self.value_empty(4, 3, 0)
528  x = self.sparse_tensor(i, v, torch.Size([3, 4, 5, 3, 0]))
529  exp_i = self.index_tensor([
530  [0, 2],
531  [0, 3],
532  [0, 4],
533  ])
534  exp_v = self.value_empty(2, 3, 0)
535  test_tensor(x, exp_i, exp_v)
536 
537  def test_clone(self):
538  def test_shape(sparse_dims, nnz, with_size):
539  x = self._gen_sparse(sparse_dims, nnz, with_size)[0]
540  if self.is_uncoalesced:
541  self.assertFalse(x.is_coalesced())
542  y = x.clone()
543  self.assertFalse(y.is_coalesced())
544  x = x.coalesce()
545  self.assertTrue(x.is_coalesced())
546  y = x.clone()
547  self.assertTrue(y.is_coalesced())
548 
549  test_shape(4, 20, 5)
550  test_shape(3, 10, [100, 100, 100, 5, 5, 5, 0])
551  test_shape(3, 0, [0, 0, 100, 5, 5, 5, 0])
552 
553  def test_Sparse_to_Sparse_copy_(self):
554  # This is for testing torch.copy_(SparseTensor, SparseTensor)
555  sparse_dims = 3
556  nnz = 10
557  sizes = [2, 3, 4, 5] # hybrid sparse
558  x1, _, _ = self._gen_sparse(sparse_dims, nnz, sizes)
559  x2, _, _ = self._gen_sparse(sparse_dims, nnz + 10, sizes)
560 
561  # test copy
562  x2_dense = x2.to_dense()
563  x1.copy_(x2)
564  self.assertEqual(x2_dense, x1.to_dense())
565 
566  # test type conversion (when x1.copy_(x2), x1.dtype should stay the same)
567  x1 = x1.to(torch.float32)
568  x2 = x2.to(torch.float64)
569  x1_dtype = x1.dtype
570  x1.copy_(x2)
571  self.assertEqual(x1_dtype, x1.dtype)
572 
573  # test no broadcast
574  self.assertRaises(RuntimeError, lambda: x1.copy_(x2.narrow_copy(0, 0, 1)))
575 
576  # test raise error on copy_() between dense and sparse Tensors
577  self.assertRaises(RuntimeError, lambda: x1.copy_(torch.randn(5, 5)))
578 
579  # test autograd
580  x1, _, _ = self._gen_sparse(sparse_dims, nnz, sizes)
581  x2, _, _ = self._gen_sparse(sparse_dims, nnz + 10, sizes)
582  x2.requires_grad_(True)
583  x1.copy_(x2)
584  y = x1 * 2
585  x2_clone = x2.clone()
586  y.backward(x2_clone)
587  expected_grad = x2_clone * 2
588  self.assertEqual(expected_grad.to_dense(), x2.grad.to_dense())
589  self.assertEqual(None, x1.grad)
590 
591  @unittest.skipIf(torch.cuda.device_count() < 2, "no multi-GPU")
592  @skipIfRocm
593  def test_Sparse_to_Sparse_copy_multi_gpu(self):
594  # This is for testing torch.copy_(SparseTensor, SparseTensor) across GPU devices
595  sparse_dims = 3
596  nnz = 10
597  sizes = [2, 3, 4, 5] # hybrid sparse
598  x1, _, _ = self._gen_sparse(sparse_dims, nnz, sizes)
599  x2, _, _ = self._gen_sparse(sparse_dims, nnz + 10, sizes)
600  x1 = x1.to('cuda:0')
601 
602  def test_cross_device(x1, x2):
603  x1_device = x1.device
604  x1.copy_(x2)
605  self.assertEqual(x2.to('cuda:0').to_dense(), x1.to_dense())
606  self.assertEqual(x1_device, x1.device)
607 
608  test_cross_device(x1, x2.to('cuda:1')) # test across gpu devices
609  test_cross_device(x1, x2.to('cpu')) # test between cpu and gpu
610 
611  # test autograd
612  x2 = x2.to('cuda:1')
613  x2.requires_grad_(True)
614  x1.copy_(x2)
615  y = x1 * 2
616  x2_clone = x2.clone().to('cuda:0')
617  y.backward(x2_clone)
618  expected_grad = x2_clone * 2
619  self.assertEqual(expected_grad.to_dense(), x2.grad.to('cuda:0').to_dense())
620  self.assertEqual(None, x1.grad)
621 
622  @cuda_only
623  def test_cuda_empty(self):
624  def test_tensor(x):
625  y = x.cuda(0)
626  self.assertEqual(x.sparse_dim(), y.sparse_dim())
627  self.assertEqual(x.dense_dim(), y.dense_dim())
628  x = y.cpu()
629  self.assertEqual(y.sparse_dim(), x.sparse_dim())
630  self.assertEqual(y.dense_dim(), x.dense_dim())
631 
632  x = torch.sparse.FloatTensor(2, 3, 4)
633  test_tensor(x)
634 
635  x = torch.sparse.FloatTensor(2, 3, 4, 0)
636  test_tensor(x)
637 
638  def test_transpose(self):
639  def test_shape(sparse_dims, nnz, with_size):
640  x = self._gen_sparse(sparse_dims, nnz, with_size)[0]
641  y = self.safeToDense(x)
642 
643  for i, j in itertools.combinations(range(4), 2):
644  x = x.transpose_(i, j)
645  y = y.transpose(i, j)
646  self.assertEqual(self.safeToDense(x), y)
647 
648  x = x.transpose(i, j)
649  y = y.transpose(i, j)
650  self.assertEqual(self.safeToDense(x), y)
651 
652  test_shape(4, 6, 3)
653  test_shape(4, 3, [7, 7, 7, 3, 3, 3, 0])
654  test_shape(4, 0, [0, 0, 7, 3, 3, 3, 0])
655 
656  @cpu_only
657  def test_coalesce_transpose_mm(self):
658  def test_shape(di, dj, dk, nnz):
659  x, _, _ = self._gen_sparse(2, nnz, [dj, di])
660  y = torch.randn(dj, dk)
661 
662  x_coalesced = x.coalesce()
663  self.assertTrue(x_coalesced.is_coalesced())
664 
665  x_coalesced_t = x_coalesced.t()
666  # Transpose is `colasced`-preserving if the indices tensor is empty.
667  self.assertEqual(x_coalesced_t.is_coalesced(), di * nnz == 0)
668 
669  res = torch.mm(x_coalesced_t, y)
670  expected = torch.mm(self.safeToDense(x_coalesced_t), y)
671  self.assertEqual(res, expected)
672 
673  test_shape(10, 20, 30, 20)
674  test_shape(0, 20, 30, 0)
675  test_shape(10, 0, 30, 0)
676  test_shape(10, 20, 0, 0)
677  test_shape(10, 20, 0, 20)
678 
679  def test_t_empty(self):
680  def test_in_place(x):
681  shape_original = x.shape
682  x.t_()
683  self.assertEqual(torch.Size([shape_original[1], shape_original[0]]), x.size())
684  self.assertEqual(0, x._indices().numel())
685  self.assertEqual(0, x._values().numel())
686  self.assertEqual(x.sparse_dim(), 2)
687  self.assertEqual(x.dense_dim(), 0)
688 
689  def test_not_in_place(x):
690  shape_original = x.shape
691  y = x.t()
692  self.assertEqual(torch.Size([shape_original[1], shape_original[0]]), y.size())
693  self.assertEqual(0, y._indices().numel())
694  self.assertEqual(0, y._values().numel())
695  self.assertEqual(x.sparse_dim(), 2)
696  self.assertEqual(x.dense_dim(), 0)
697 
698  x = self.sparse_empty(2, 3)
699  test_in_place(x)
700  test_not_in_place(x)
701 
702  x = self.sparse_empty(2, 0)
703  test_in_place(x)
704  test_not_in_place(x)
705 
706  def test_add_zeros(self):
707  def test_shape(sparse_dims, nnz, sizes):
708  x, _, _ = self._gen_sparse(sparse_dims, nnz, sizes)
709  zeros = torch.zeros(sizes, layout=torch.sparse_coo).to(x.device)
710  r1 = zeros + x
711  r2 = x + zeros
712  self.assertEqual(r1, x)
713  self.assertEqual(r2, x)
714 
715  test_shape(1, 20, [1])
716  test_shape(4, 20, [3, 17, 19, 5])
717  test_shape(2, 20, [3, 17, 19, 5])
718  test_shape(2, 20, [3, 17, 19, 0])
719 
720  def test_cat(self):
721  # shapes: list of tuples (sparse_dims, nnz, sizes)
722  def test_shapes(shapes, dim, fail_message=None):
723  inputs = [self._gen_sparse(shape[0], shape[1], shape[2])[0]
724  for shape in shapes]
725  if fail_message:
726  with self.assertRaisesRegex(RuntimeError, fail_message):
727  torch.cat(inputs, dim)
728  else:
729  result = torch.cat(inputs, dim)
730  dense_result = torch.cat([t.to_dense() for t in inputs], dim)
731  self.assertEqual(dense_result, result.to_dense())
732 
733  test_shapes(
734  [(3, 10, [2, 3, 4]), (3, 10, [2, 1, 4]), (3, 10, [2, 4, 4])], 1)
735 
736  # mismatched sizes
737  test_shapes([(3, 10, [2, 3, 4]), (3, 10, [2, 1, 4])], 0,
738  "All tensors must have the same shape: \\[2, 3, 4].*\\[2, 1, 4]")
739  # hybrid sparse/dense
740  test_shapes(
741  [(2, 10, [2, 3, 4]), (2, 10, [2, 1, 4]), (2, 10, [2, 4, 4])], 1)
742  # cat along dense dim
743  test_shapes([(2, 10, [2, 3, 4]), (2, 10, [2, 3, 7])], 2)
744  test_shapes([(1, 10, [2, 3, 4]), (1, 10, [2, 3, 4])], 1)
745  test_shapes([(1, 10, [2, 3, 4]), (1, 10, [2, 3, 4])], 2)
746  # mismatched dimensions
747  test_shapes([(2, 10, [2, 3, 4]), (3, 10, [2, 3, 4])], 0,
748  "All tensors must have the same.*2, 1, but tensor at position 1 has 3, 0.")
749  # wrapped dimension
750  test_shapes(
751  [(3, 10, [2, 3, 4]), (3, 10, [2, 1, 4]), (3, 10, [2, 4, 4])], -2)
752 
753  # sparse with dense
754  sp = self._gen_sparse(3, 10, [2, 3, 4])[0]
755  dn = sp.to_dense()
756  with self.assertRaisesRegex(RuntimeError,
757  "Concatenating sparse tensors, but a dense tensor was found at position 1."):
758  torch.cat((sp, dn))
759 
760  def test_unsqueeze(self):
761  def test_shape(sparse_dims, nnz, sizes, unsqueeze_dim, fail_message=None):
762  x, _, _ = self._gen_sparse(sparse_dims, nnz, sizes)
763  if fail_message:
764  with self.assertRaisesRegex(IndexError, fail_message):
765  torch.unsqueeze(x, unsqueeze_dim)
766  else:
767  result = torch.unsqueeze(x, unsqueeze_dim)
768  dense_result = torch.unsqueeze(x.to_dense(), unsqueeze_dim)
769  self.assertEqual(dense_result, result.to_dense())
770 
771  # basic case
772  test_shape(3, 10, [5, 7, 11], 0)
773 
774  # hybrid sparse/dense, unsqueeze along sparse dim
775  test_shape(3, 10, [5, 7, 11, 13, 17], 0)
776  test_shape(3, 10, [5, 7, 11, 13, 17], 3)
777 
778  # unsqueeze along dense dimensions
779  test_shape(3, 10, [5, 7, 11, 13, 17], 4)
780  test_shape(3, 10, [5, 7, 11, 13, 17], 5)
781 
782  # wrapped dimensions
783  test_shape(3, 10, [5, 7, 11, 13, 17], -1)
784  test_shape(3, 10, [5, 7, 11, 13, 17], -6)
785 
786  # bounds
787  test_shape(3, 10, [5, 7, 11, 13, 17], -7, "Dimension out of range")
788  test_shape(3, 10, [5, 7, 11, 13, 17], 6, "Dimension out of range")
789 
790  @cpu_only
791  def test_mm(self):
792  def test_shape(di, dj, dk, nnz):
793  x, _, _ = self._gen_sparse(2, nnz, [di, dj])
794  t = torch.randn(di, dk)
795  y = torch.randn(dj, dk)
796  alpha = random.random()
797  beta = random.random()
798 
799  res = torch.addmm(alpha, t, beta, x, y)
800  expected = torch.addmm(alpha, t, beta, self.safeToDense(x), y)
801  self.assertEqual(res, expected)
802 
803  res = torch.addmm(t, x, y)
804  expected = torch.addmm(t, self.safeToDense(x), y)
805  self.assertEqual(res, expected)
806 
807  res = torch.mm(x, y)
808  expected = torch.mm(self.safeToDense(x), y)
809  self.assertEqual(res, expected)
810 
811  test_shape(10, 100, 100, 20)
812  test_shape(100, 1000, 200, 20)
813  test_shape(64, 10000, 300, 20)
814  test_shape(0, 100, 100, 0)
815  test_shape(10, 0, 100, 0)
816  test_shape(10, 100, 0, 0)
817  test_shape(10, 100, 0, 20)
818 
819  @cpu_only
820  def test_saddmm(self):
821  def test_shape(di, dj, dk, nnz):
822  x = self._gen_sparse(2, nnz, [di, dj])[0]
823  t = self._gen_sparse(2, nnz, [di, dk])[0]
824  y = torch.randn(dj, dk)
825  alpha = random.random()
826  beta = random.random()
827 
828  res = torch.saddmm(alpha, t, beta, x, y)
829  expected = torch.addmm(alpha, self.safeToDense(t), beta, self.safeToDense(x), y)
830  self.assertEqual(self.safeToDense(res), expected)
831 
832  res = torch.saddmm(t, x, y)
833  expected = torch.addmm(self.safeToDense(t), self.safeToDense(x), y)
834  self.assertEqual(self.safeToDense(res), expected)
835 
836  res = torch.smm(x, y)
837  expected = torch.mm(self.safeToDense(x), y)
838  self.assertEqual(self.safeToDense(res), expected)
839 
840  test_shape(7, 5, 3, 20)
841  test_shape(1000, 100, 100, 20)
842  test_shape(3000, 64, 300, 20)
843  test_shape(0, 100, 100, 0)
844  test_shape(1000, 0, 100, 0)
845  test_shape(1000, 100, 0, 0)
846 
847  def test_sparse_addmm(self):
848  def test_shape(m, n, p, nnz, broadcast):
849  if broadcast:
850  D1 = torch.randn((), device=self.device).requires_grad_(True)
851  else:
852  D1 = torch.randn(n, p, device=self.device).requires_grad_(True)
853  D2 = torch.randn(m, p, device=self.device).requires_grad_(True)
854  S = self._gen_sparse(2, nnz, [n, m])[0]
855  S_dense = S.to_dense().requires_grad_(True)
856  S.requires_grad_(True)
857  self.assertEqual(torch.sparse.addmm(D1, S, D2), torch.addmm(D1, S_dense, D2))
858 
859  def fn(S, D1, D2):
860  return torch.sparse.addmm(D1, S, D2)
861  gradcheck(fn, (S, D1, D2), check_sparse_nnz=True)
862 
863  test_shape(7, 8, 9, 20, False)
864  test_shape(7, 8, 9, 20, True)
865 
866  def test_sparse_mm(self):
867  def test_shape(d1, d2, d3, nnz):
868  D = torch.randn(d2, d3, device=self.device).requires_grad_(True)
869  S = self._gen_sparse(2, nnz, [d1, d2])[0]
870  S_dense = S.to_dense().requires_grad_(True)
871  S.requires_grad_(True)
872  self.assertEqual(torch.sparse.mm(S, D), torch.mm(S_dense, D))
873 
874  def fn(S, D):
875  return torch.sparse.mm(S, D)
876  gradcheck(fn, (S, D), check_sparse_nnz=True)
877 
878  test_shape(7, 8, 9, 20)
879 
880  @skipIfRocm
881  def test_dsmm(self):
882  def test_shape(di, dj, dk, nnz):
883  x = self._gen_sparse(2, nnz, [di, dj])[0]
884  y = self.randn(dj, dk)
885 
886  res = torch.dsmm(x, y)
887  expected = torch.mm(self.safeToDense(x), y)
888  self.assertEqual(res, expected)
889 
890  test_shape(7, 5, 3, 20)
891  test_shape(1000, 100, 100, 20)
892  test_shape(3000, 64, 300, 20)
893  test_shape(0, 100, 100, 0)
894  test_shape(1000, 0, 100, 0)
895  test_shape(1000, 100, 0, 0)
896  test_shape(1000, 100, 0, 20)
897 
898  @skipIfRocm
899  def test_hsmm(self):
900  def test_shape(di, dj, dk, nnz):
901  x = self._gen_sparse(2, nnz, [di, dj])[0]
902  y = self.randn(dj, dk)
903 
904  res = torch.hsmm(x, y)
905  expected = torch.mm(self.safeToDense(x), y)
906  self.assertEqual(res.to_dense(), expected)
907 
908  test_shape(7, 5, 3, 20)
909  test_shape(1000, 100, 100, 20)
910  test_shape(3000, 64, 300, 20)
911  test_shape(0, 100, 100, 0)
912  test_shape(1000, 0, 100, 0)
913  test_shape(1000, 100, 0, 0)
914  test_shape(1000, 100, 0, 20)
915 
916  def _test_spadd_shape(self, nnz, shape_i, shape_v=None):
917  shape = shape_i + (shape_v or [])
918  x, _, _ = self._gen_sparse(len(shape_i), nnz, shape)
919  y = self.randn(*shape)
920  r = random.random()
921 
922  res = torch.add(y, r, x)
923  expected = y + r * self.safeToDense(x)
924 
925  self.assertEqual(res, expected)
926 
927  # Non contiguous dense tensor
928  s = list(shape)
929  s[0] = shape[-1]
930  s[-1] = shape[0]
931  y = self.randn(*s)
932  y.transpose_(0, len(s) - 1)
933  r = random.random()
934 
935  res = torch.add(y, r, x)
936  expected = y + r * self.safeToDense(x)
937 
938  self.assertEqual(res, expected)
939 
940  x, i, v = self._gen_sparse(len(shape_i), nnz, shape)
941  nnz = i.size(1)
942 
943  # Non contiguous sparse indices tensor
944  x_ = self.sparse_tensor(i[:, ::2], v[:int(nnz / 2)], x.shape)
945  res = torch.add(y, r, x_)
946  expected = y + r * self.safeToDense(x_)
947  self.assertEqual(res, expected)
948 
949  # Non contiguous sparse values tensor
950  x_ = self.sparse_tensor(i[:, :int(nnz / 2)], v[::2], x.shape)
951  res = torch.add(y, r, x_)
952  expected = y + r * self.safeToDense(x_)
953  self.assertEqual(res, expected)
954 
955  # Non contiguous sparse indices and values tensors
956  x_ = self.sparse_tensor(i[:, 1::2], v[1::2], x.shape)
957  res = torch.add(y, r, x_)
958  expected = y + r * self.safeToDense(x_)
959  self.assertEqual(res, expected)
960 
961  def test_spadd(self):
962  self._test_spadd_shape(10, [5, 6])
963  self._test_spadd_shape(10, [10, 10, 10])
964  self._test_spadd_shape(10, [50, 30, 20])
965  self._test_spadd_shape(10, [5, 5, 5, 5, 5, 5])
966  self._test_spadd_shape(0, [0, 30, 20])
967  self._test_spadd_shape(0, [50, 0, 20])
968  self._test_spadd_shape(0, [50, 30, 0])
969 
970  def test_spadd_hybrid(self):
971  self._test_spadd_shape(10, [5, 6], [2, 3])
972  self._test_spadd_shape(10, [10, 10, 10], [3])
973  self._test_spadd_shape(10, [50, 30, 20], [2])
974  self._test_spadd_shape(10, [5, 5, 5, 5, 5, 5], [2])
975  self._test_spadd_shape(0, [0, 30, 20], [2, 0])
976  self._test_spadd_shape(0, [50, 0, 20], [2, 0])
977  self._test_spadd_shape(0, [50, 30, 0], [2, 0])
978  self._test_spadd_shape(10, [50, 30, 20], [2, 0])
979 
980  def test_norm(self):
981  def test_shape(sparse_dims, nnz, with_size):
982  x, _, _ = self._gen_sparse(sparse_dims, nnz, with_size)
983  y = x.coalesce()
984  self.assertEqual(x.norm(), y._values().norm())
985 
986  test_shape(3, 10, 100)
987  test_shape(4, 10, [100, 100, 100, 5, 5, 5, 0])
988  test_shape(4, 0, [0, 0, 100, 5, 5, 5, 0])
989 
990  @skipIfRocm
991  def test_sparse_sum(self):
992 
993  def run_tests(S, td=None):
994  D = S.coalesce().to_dense().detach().requires_grad_(True)
995  mask = (D == 0)
996  if td is None:
997  S_sum = torch.sparse.sum(S)
998  D_sum = D.sum()
999  self.assertEqual(S_sum, D_sum)
1000 
1001  def fn(S):
1002  res = torch.sparse.sum(S)
1003  if res.is_sparse:
1004  res = res.to_dense()
1005  return res
1006  gradcheck(fn, (S,), check_sparse_nnz=True)
1007 
1008  else:
1009  S_sum = torch.sparse.sum(S, td)
1010  D_sum = D.sum(td)
1011  self.assertEqual(S_sum.to_dense() if S_sum.is_sparse else S_sum, D_sum)
1012 
1013  def fn(S):
1014  res = torch.sparse.sum(S, td)
1015  if res.is_sparse:
1016  res = res.to_dense()
1017  return res
1018  gradcheck(fn, (S,), check_sparse_nnz=True)
1019 
1020  nnz = 10
1021  sparse_dims = 2
1022  with_size = [5, 5, 1, 4] # use a dense dim = 1 to test for squeeze
1023  test_dims = []
1024  for i in range(1, 5):
1025  test_dims += itertools.combinations(range(len(with_size)), i)
1026 
1027  # https://github.com/pytorch/pytorch/issues/16501
1028  x = torch.tensor([[1., 0., 0., 1.],
1029  [0., 1., 0., 0.],
1030  [0., 1., 1., 0.],
1031  [0., 1., 0., 2.]]).to_sparse()
1032  self.assertEqual(torch.sparse.sum(x, dim=0), torch.sparse.sum(x, dim=-2))
1033  self.assertEqual(torch.sum(x.to_dense(), dim=0), torch.sparse.sum(x, dim=0).to_dense())
1034 
1035  # not support SparseTensor.sum()
1036  S = self._gen_sparse(sparse_dims, nnz, with_size)[0]
1037  self.assertRaises(RuntimeError, lambda: S.sum())
1038 
1039  # dim out of range
1040  self.assertRaises(IndexError, lambda: torch.sparse.sum(S, 5))
1041 
1042  # dim 0 appears multiple times in the list of dims
1043  self.assertRaises(RuntimeError, lambda: torch.sparse.sum(S, [0, 0]))
1044 
1045  # sum an empty tensor
1046  empty_S = torch.sparse_coo_tensor(size=with_size)
1047  self.assertRaises(RuntimeError, lambda: torch.sparse.sum(empty_S, [0]))
1048  self.assertEqual(torch.sparse.sum(empty_S), torch.tensor(0,))
1049  empty_S.requires_grad_(True)
1050  empty_S_sum = torch.sparse.sum(empty_S)
1051  empty_S_sum.backward()
1052  self.assertEqual(empty_S.grad.to_dense(), empty_S.clone().detach().to_dense())
1053 
1054  # test values().sum()
1055  S = self._gen_sparse(sparse_dims, nnz, with_size)[0]
1056  run_tests(S.requires_grad_(True))
1057 
1058  for test_dim in test_dims:
1059  S = self._gen_sparse(sparse_dims, nnz, with_size)[0]
1060  run_tests(S.requires_grad_(True), test_dim)
1061 
1062  def _test_basic_ops_shape(self, nnz_x1, nnz_x2, shape_i, shape_v=None):
1063  shape = shape_i + (shape_v or [])
1064  x1, _, _ = self._gen_sparse(len(shape_i), nnz_x1, shape)
1065  x2, _, _ = self._gen_sparse(len(shape_i), nnz_x2, shape)
1066 
1067  y1 = x1 + x2
1068  y2 = x1.clone()
1069  y2.add_(x2)
1070  expected = self.safeToDense(x1) + self.safeToDense(x2)
1071  self.assertEqual(self.safeToDense(y1), expected)
1072  self.assertEqual(self.safeToDense(y2), expected)
1073 
1074  y1 = x1 - x2
1075  y2 = x1.clone()
1076  y2.sub_(x2)
1077  expected = self.safeToDense(x1) - self.safeToDense(x2)
1078  self.assertEqual(self.safeToDense(y1), expected)
1079  self.assertEqual(self.safeToDense(y2), expected)
1080 
1081  y1 = x1 * x2
1082  y2 = x1.clone()
1083  y2.mul_(x2)
1084  expected = self.safeToDense(x1) * self.safeToDense(x2)
1085  self.assertEqual(self.safeToDense(y1), expected)
1086  self.assertEqual(self.safeToDense(y2), expected)
1087 
1088  y1 = x1 * 37.5
1089  y2 = x1.clone()
1090  y2.mul_(37.5)
1091  expected = self.safeToDense(x1) * 37.5
1092  self.assertEqual(self.safeToDense(y1), expected)
1093  self.assertEqual(self.safeToDense(y2), expected)
1094 
1095  y1 = x1 / 37.5
1096  y2 = x1.clone()
1097  y2.div_(37.5)
1098  expected = self.safeToDense(x1) / 37.5
1099  self.assertEqual(self.safeToDense(y1), expected)
1100  self.assertEqual(self.safeToDense(y2), expected)
1101 
1102  # TODO: add back inplace support
1103  y1 = x1 ** 2
1104  y2 = x1.clone()
1105  y2 = y2.pow(2)
1106  expected = self.safeToDense(x1) ** 2
1107  self.assertEqual(self.safeToDense(y1), expected)
1108  self.assertEqual(self.safeToDense(y2), expected)
1109 
1110  y = x1.clone()
1111  y.zero_()
1112  expected = torch.zeros(x1.size())
1113  self.assertEqual(self.safeToDense(y), expected)
1114 
1115  self.assertEqual(x1.is_coalesced(), not self.is_uncoalesced)
1116  y = x1.coalesce()
1117  z = x1.coalesce()
1118  self.assertEqual(x1.is_coalesced(), not self.is_uncoalesced)
1119  self.assertTrue(y.is_coalesced())
1120  self.assertEqual(x1, y)
1121  y._values().add_(1)
1122  if not x1.is_coalesced():
1123  # check that coalesce is out of place if the original tensor is not
1124  # coalesced.
1125  self.assertEqual(z._values() + 1, y._values())
1126  else:
1127  # check that coalesce is in-place if the original tensor is
1128  # coalesced.
1129  self.assertEqual(z._values(), y._values())
1130 
1131  def test_basic_ops(self):
1132  self._test_basic_ops_shape(9, 12, [5, 6])
1133  self._test_basic_ops_shape(9, 12, [10, 10, 10])
1134  self._test_basic_ops_shape(9, 12, [50, 30, 20])
1135  self._test_basic_ops_shape(9, 12, [5, 5, 5, 5, 5, 5])
1136  self._test_basic_ops_shape(0, 12, [10, 10, 10])
1137  self._test_basic_ops_shape(9, 0, [10, 10, 10])
1138  self._test_basic_ops_shape(0, 0, [10, 10, 10])
1139  self._test_basic_ops_shape(0, 0, [10, 10, 0])
1140 
1141  @skipIfRocm
1142  def test_basic_ops_hybrid(self):
1143  self._test_basic_ops_shape(9, 12, [5, 6], [2, 3])
1144  self._test_basic_ops_shape(9, 12, [10, 10, 10], [3])
1145  self._test_basic_ops_shape(9, 12, [50, 30, 20], [2])
1146  self._test_basic_ops_shape(9, 12, [5, 5, 5, 5, 5, 5], [2])
1147  self._test_basic_ops_shape(0, 12, [10, 10, 10], [2])
1148  self._test_basic_ops_shape(9, 0, [10, 10, 10], [2])
1149  self._test_basic_ops_shape(0, 0, [10, 10, 10], [2])
1150  self._test_basic_ops_shape(9, 12, [10, 10, 10], [2, 0])
1151  self._test_basic_ops_shape(0, 12, [10, 10, 10], [2, 0])
1152  self._test_basic_ops_shape(9, 0, [10, 10, 10], [2, 0])
1153  self._test_basic_ops_shape(0, 0, [10, 10, 10], [2, 0])
1154  self._test_basic_ops_shape(0, 0, [10, 10, 0], [2, 0])
1155 
1156  def test_add_dense_sparse_mismatch(self):
1157  def test_shape(dense_size, sparse_dims_shape, dense_dims_shape, sparse_size):
1158  x = torch.zeros(dense_size, dtype=self.value_dtype, device=self.device)
1159  sparse_y = self.sparse_tensor(torch.zeros(sparse_dims_shape, dtype=torch.int64, device=self.device),
1160  torch.randn(dense_dims_shape, dtype=self.value_dtype, device=self.device),
1161  torch.Size(sparse_size))
1162  with self.assertRaisesRegex(
1163  RuntimeError,
1164  "add: expected 'self' and 'other' to have same size"):
1165  x + sparse_y
1166 
1167  test_shape([3, 4], [1, 4], [4, 4, 4], [3, 4, 4])
1168  test_shape([3, 4, 0], [1, 4], [4, 4, 4, 0], [3, 4, 4, 0])
1169 
1170  def _test_sparse_mask_shape(self, nnz_x1, nnz_x2, shape_i, shape_v=None):
1171  shape = shape_i + (shape_v or [])
1172  x1, _, _ = self._gen_sparse(len(shape_i), nnz_x1, shape)
1173  x2, _, _ = self._gen_sparse(len(shape_i), nnz_x2, shape)
1174 
1175  y1 = x1 + x2
1176  y2 = x1.clone()
1177  y2.add_(x2)
1178  expected = self.safeToDense(x1) + self.safeToDense(x2)
1179  self.assertEqual(self.safeToDense(y1), expected)
1180  self.assertEqual(self.safeToDense(y2), expected)
1181 
1182  def _test_sparse_mask_fixed(self):
1183  i = self.index_tensor([
1184  [1, 3, 0, 4],
1185  [2, 1, 2, 3],
1186  ])
1187  v = self.value_tensor([1, 2, 3, 4])
1188  x = self.sparse_tensor(i, v, torch.Size([5, 4])).coalesce()
1189  dense = self.value_tensor([
1190  [1, 2, 3, 4],
1191  [5, 6, 7, 8],
1192  [9, 10, 11, 12],
1193  [13, 14, 15, 16],
1194  [17, 18, 19, 20],
1195  ])
1196  exp_v = self.value_tensor([7, 14, 3, 20])
1197  res = dense.sparse_mask(x)
1198  expected = self.sparse_tensor(i, exp_v, torch.Size([5, 4]))
1199  self.assertEqual(res, expected)
1200 
1201  i = self.index_tensor([
1202  [1, 3, 0, 4],
1203  [2, 1, 2, 3],
1204  ])
1205  v = self.value_empty(4, 0)
1206  x = self.sparse_tensor(i, v, torch.Size([5, 4, 0])).coalesce()
1207  dense = self.value_empty(5, 4, 0)
1208  exp_v = self.value_empty(4, 0)
1209  res = dense.sparse_mask(x)
1210  expected = self.sparse_tensor(i, exp_v, torch.Size([5, 4, 0]))
1211  self.assertEqual(res, expected)
1212 
1213  def test_sparse_mask(self):
1215 
1216  self._test_sparse_mask_shape(9, 12, [5, 6])
1217  self._test_sparse_mask_shape(9, 12, [10, 10, 10])
1218  self._test_sparse_mask_shape(9, 12, [50, 30, 20])
1219  self._test_sparse_mask_shape(9, 12, [5, 5, 5, 5, 5, 5])
1220  self._test_sparse_mask_shape(0, 12, [10, 10, 10])
1221  self._test_sparse_mask_shape(9, 0, [10, 10, 10])
1222  self._test_sparse_mask_shape(0, 0, [10, 10, 10])
1223  self._test_sparse_mask_shape(0, 0, [10, 10, 0])
1224 
1225  def _test_sparse_mask_hybrid_fixed(self):
1226  i = self.index_tensor([
1227  [1, 3, 0, 4],
1228  [2, 1, 2, 3],
1229  ])
1230  v = self.value_tensor([[1, 2], [2, 3], [3, 4], [4, 5]])
1231  # TODO: This is also testing that, if coalesce is a no-op,
1232  # the indices don't get permuted. I don't know if we actually
1233  # want to give this invariant.
1234  x = self.sparse_tensor(i, v, torch.Size([5, 4, 2])).coalesce()
1235  dense = self.value_tensor([
1236  [[1, 3], [2, 2], [3, 3], [4, 2]],
1237  [[5, 7], [6, 7], [7, 9], [8, 9]],
1238  [[9, 2], [10, 4], [11, 1], [12, 3]],
1239  [[13, 5], [14, 1], [15, 1], [16, 6]],
1240  [[17, 7], [18, 2], [19, 7], [20, 1]],
1241  ])
1242  res = dense.sparse_mask(x)
1243  exp_v = self.value_tensor([[7, 9], [14, 1], [3, 3], [20, 1]])
1244  expected = self.sparse_tensor(i, exp_v, torch.Size([5, 4, 2]))
1245  self.assertEqual(res, expected)
1246 
1247  i = self.index_tensor([
1248  [1, 3, 0, 4],
1249  [2, 1, 2, 3],
1250  ])
1251  v = self.value_empty(4, 2, 0)
1252  x = self.sparse_tensor(i, v, torch.Size([5, 4, 2, 0])).coalesce()
1253  dense = self.value_empty(5, 4, 2, 0)
1254  res = dense.sparse_mask(x)
1255  exp_v = self.value_empty(4, 2, 0)
1256  expected = self.sparse_tensor(i, exp_v, torch.Size([5, 4, 2, 0]))
1257  self.assertEqual(res, expected)
1258 
1259  def test_sparse_mask_hybrid(self):
1261 
1262  self._test_sparse_mask_shape(9, 12, [5, 6], [2, 3])
1263  self._test_sparse_mask_shape(9, 12, [10, 10, 10], [3])
1264  self._test_sparse_mask_shape(9, 12, [50, 30, 20], [2])
1265  self._test_sparse_mask_shape(9, 12, [5, 5, 5, 5, 5, 5], [2])
1266  self._test_sparse_mask_shape(0, 12, [10, 10, 10], [2])
1267  self._test_sparse_mask_shape(9, 0, [10, 10, 10], [2])
1268  self._test_sparse_mask_shape(0, 0, [10, 10, 10], [2])
1269  self._test_sparse_mask_shape(9, 12, [10, 10, 10], [2, 0])
1270  self._test_sparse_mask_shape(0, 12, [10, 10, 10], [2, 0])
1271  self._test_sparse_mask_shape(9, 0, [10, 10, 10], [2, 0])
1272  self._test_sparse_mask_shape(0, 0, [10, 10, 10], [2, 0])
1273  self._test_sparse_mask_shape(0, 0, [10, 10, 0], [2, 0])
1274 
1275  def _test_zeros(self, nnzs, shape, out_shape_i, out_shape_v=None):
1276  out_shape = out_shape_i + (out_shape_v or [])
1277  for nnz in nnzs:
1278  out, _, _ = self._gen_sparse(len(out_shape_i), nnz, out_shape)
1279  torch.zeros(*shape, out=out)
1280  self.assertEqual(tuple(out.size()), tuple(shape))
1281  self.assertTrue(out._indices().numel() == out._values().numel() == 0)
1282  self.assertEqual(out._nnz(), 0)
1283  self.assertEqual(out.sparse_dim(), len(shape))
1284  self.assertEqual(out.dense_dim(), 0)
1285 
1286  def test_zeros(self):
1287  def test_shape(i_shapes, v_shapes, shape, nnzs):
1288  for i_dim in range(1, len(i_shapes) + 1):
1289  for v_dim in range(len(v_shapes) + 1):
1290  self._test_zeros(nnzs, shape, i_shapes[:i_dim], v_shapes[:v_dim])
1291  test_shape([2, 3, 4], [3, 4, 5, 6], [2, 3, 4], [9, 12])
1292  test_shape([0, 3, 4], [3, 4, 5, 6], [2, 3, 4], [0])
1293  test_shape([2, 3, 4], [0, 4, 5, 6], [2, 3, 4], [9, 12])
1294  test_shape([2, 3, 4], [3, 4, 5, 6], [2, 3, 0], [9, 12])
1295  test_shape([0, 3, 4], [3, 4, 5, 6], [2, 3, 0], [0])
1296  test_shape([2, 3, 4], [0, 4, 5, 6], [2, 3, 0], [9, 12])
1297 
1298  def _test_zeros_like(self, nnzs, template_shape_i, template_shape_v=None):
1299  template_shape_v = template_shape_v or []
1300  template_shape = template_shape_i + template_shape_v
1301  for nnz in nnzs:
1302  t, _, _ = self._gen_sparse(len(template_shape_i), nnz, template_shape)
1303  res = torch.zeros_like(t)
1304  self.assertEqual(tuple(res.size()), tuple(template_shape))
1305  self.assertTrue(res._indices().numel() == res._values().numel() == 0)
1306  self.assertEqual(res._nnz(), 0)
1307  self.assertEqual(res.sparse_dim(), len(template_shape_i))
1308  self.assertEqual(res.dense_dim(), len(template_shape_v))
1309 
1310  @skipIfRocm
1311  def test_zeros_like(self):
1312  def test_shape(i_shapes, v_shapes, nnzs):
1313  for i_dim in range(1, len(i_shapes) + 1):
1314  for v_dim in range(len(v_shapes) + 1):
1315  self._test_zeros_like(nnzs, i_shapes[:i_dim], v_shapes[:v_dim])
1316  test_shape([2, 3, 4], [3, 4, 5, 6], [9, 12])
1317  test_shape([0, 3, 4], [3, 4, 5, 6], [0])
1318  test_shape([2, 3, 4], [0, 4, 5, 6], [9, 12])
1319  test_shape([2, 3, 4], [3, 4, 5, 6], [9, 12])
1320  test_shape([0, 3, 4], [3, 4, 5, 6], [0])
1321  test_shape([2, 3, 4], [0, 4, 5, 6], [9, 12])
1322 
1323  def _test_narrow(self, input, narrow_args):
1324  expected = input.to_dense().narrow(*narrow_args)
1325  self.assertEqual(expected, input.narrow_copy(*narrow_args).to_dense())
1326 
1327  def _all_narrow_combs(self, shape):
1328  for dim, dim_sz in enumerate(shape):
1329  for start in range(dim_sz):
1330  for length in range(dim_sz - start):
1331  yield [dim, start, length]
1332 
1333  def test_narrow(self):
1334  shape = [3, 3, 4, 2]
1335  input, _, _ = self._gen_sparse(4, 19, shape)
1336  for narrow_args in self._all_narrow_combs(shape):
1337  self._test_narrow(input, narrow_args)
1338 
1339  self.assertRaises(RuntimeError, lambda: input.narrow_copy(-1, 0, 3)) # dim < 0
1340  self.assertRaises(RuntimeError, lambda: input.narrow_copy(10, 0, 3)) # dim > input.dim()
1341  self.assertRaises(RuntimeError, lambda: input.narrow_copy(0, shape[0] + 1, 3)) # start > size of dim
1342  self.assertRaises(RuntimeError, lambda: input.narrow_copy(0, 2, shape[0])) # start+length > size of dim
1343 
1344  with_dense, _, _ = self._gen_sparse(2, 7, shape)
1345  for narrow_args in self._all_narrow_combs(shape):
1346  self._test_narrow(with_dense, narrow_args)
1347 
1348  self.assertRaises(RuntimeError, lambda: with_dense.narrow_copy(10, 0, 3)) # dim > sparseDim + denseDim
1349 
1350  def _test_log1p_tensor(self, input, dense_tensor):
1351  expected_output = dense_tensor.log1p()
1352  self.assertEqual(expected_output, input.log1p().to_dense())
1353  self.assertEqual(expected_output, input.coalesce().log1p_().to_dense())
1354 
1355  # test in-place op on uncoalesced input
1356  with self.assertRaisesRegex(RuntimeError, "in-place on uncoalesced tensors is not supported yet"):
1357  input.log1p_()
1358 
1359  input.requires_grad_()
1360  self.assertTrue(input.requires_grad)
1361 
1362  # test autograd
1363  x = input.clone()
1364  y = input.log1p()
1365  with self.assertRaisesRegex(RuntimeError, "log1p of a sparse tensor is made to be non-differentiable"):
1366  y.backward(x)
1367 
1368  def test_log1p(self):
1369  input = torch.sparse_coo_tensor(
1370  torch.LongTensor([[0], [1], [2]]).transpose(1, 0).clone().detach(),
1371  torch.FloatTensor([3, 4, 5]),
1372  torch.Size([3]),
1373  device=self.device)
1374  self._test_log1p_tensor(input, torch.as_tensor([3., 4., 5.]))
1375 
1376  # test uncoalesced input
1377  input_uncoalesced = torch.sparse_coo_tensor(
1378  torch.LongTensor([[0], [1], [2], [0], [1], [2]]).transpose(1, 0).clone().detach(),
1379  torch.FloatTensor([2, 3, 4, 1, 1, 1]),
1380  torch.Size([3]),
1381  device=self.device)
1382  self._test_log1p_tensor(input_uncoalesced, torch.as_tensor([3., 4., 5.]))
1383 
1384  input = torch.sparse_coo_tensor(
1385  torch.zeros([2, 0]),
1386  torch.zeros([0, 5, 5, 5, 5, 5, 5, 0]),
1387  torch.Size([0, 0, 5, 5, 5, 5, 5, 5, 0]),
1388  device=self.device)
1389  self._test_log1p_tensor(input, torch.zeros([0, 0, 5, 5, 5, 5, 5, 5, 0]))
1390 
1391  input = torch.sparse_coo_tensor(
1392  torch.zeros([1, 5]),
1393  torch.zeros([5, 6, 0]),
1394  torch.Size([5, 6, 0]),
1395  device=self.device)
1396  self._test_log1p_tensor(input, torch.zeros([5, 6, 0]))
1397 
1398  @skipIfRocm
1399  def test_sparse_add_coalesce(self):
1400  i = self.index_tensor([[1, 2, 1]])
1401  v = self.value_tensor([3, 4, 5])
1402  x = self.sparse_tensor(i, v, torch.Size([3]))
1403  y = self.sparse_tensor(i, v, torch.Size([3]))
1404  z = x + y
1405 
1406  self.assertFalse(z._indices().numel() != 2 and z.is_coalesced())
1407 
1408  i = self.index_tensor([[1, 2, 1]])
1409  v = self.value_empty(3, 0)
1410  x = self.sparse_tensor(i, v, torch.Size([3, 0]))
1411  y = self.sparse_tensor(i, v, torch.Size([3, 0]))
1412  z = x + y
1413 
1414  self.assertFalse(z._indices().numel() != 2 and z.is_coalesced())
1415 
1416  @cuda_only
1417  def test_storage_not_null(self):
1418  x = torch.cuda.sparse.FloatTensor(2)
1419  self.assertNotEqual(x.get_device(), -1)
1420 
1421  x = torch.cuda.sparse.FloatTensor(2, 0)
1422  self.assertNotEqual(x.get_device(), -1)
1423 
1424  @cuda_only
1425  @unittest.skipIf(torch.cuda.device_count() < 2, "only one GPU detected")
1426  def test_same_gpu(self):
1427  def check_device(x, device_id):
1428  self.assertEqual(x.get_device(), device_id)
1429  self.assertEqual(x._values().get_device(), device_id)
1430  self.assertEqual(x._indices().get_device(), device_id)
1431 
1432  i = self.index_tensor([[2]]).cuda(1)
1433  v = self.value_tensor([5]).cuda(1)
1434  x = self.sparse_tensor(i, v, torch.Size([3]), device=1)
1435  check_device(x, 1)
1436 
1437  i = self.index_tensor([[2]]).cuda(1)
1438  v = self.value_empty(1, 0).cuda(1)
1439  x = self.sparse_tensor(i, v, torch.Size([3, 0]), device=1)
1440  check_device(x, 1)
1441 
1442  x = self.sparse_empty(3, device=1)
1443  check_device(x, 1)
1444 
1445  x = self.sparse_empty(3, 0, device=1)
1446  check_device(x, 1)
1447 
1448  i = self.index_tensor([[2]]).cuda(1)
1449  v = self.value_tensor([5]).cuda(0)
1450  # NB: non-legacy constructor allows this and moves indices
1451  self.assertRaises(RuntimeError, lambda: self.legacy_sparse_tensor(i, v, torch.Size([3])))
1452 
1453  i = self.index_tensor([[2]]).cuda(1)
1454  v = self.value_empty(1, 0).cuda(0)
1455  # NB: non-legacy constructor allows this and moves indices
1456  self.assertRaises(RuntimeError, lambda: self.legacy_sparse_tensor(i, v, torch.Size([3, 0])))
1457 
1458  def _test_new_device(self, size, device):
1459  with torch.cuda.device(device):
1460  x = torch.cuda.sparse.DoubleTensor(*size)
1461  self.assertEqual(x.get_device(), device)
1462  x1 = x.new()
1463  x2 = x.new(2, 3)
1464  self.assertEqual(x1.get_device(), device)
1465  self.assertEqual(x2.get_device(), device)
1466 
1467  @cuda_only
1468  def test_new_device_single_gpu(self):
1469  self._test_new_device((), 0)
1470  self._test_new_device((30, 20), 0)
1471  self._test_new_device((30, 20, 10), 0)
1472  self._test_new_device((30, 20, 10, 0), 0)
1473 
1474  @cuda_only
1475  @unittest.skipIf(torch.cuda.device_count() < 2, "only one GPU detected")
1476  def test_new_device_multi_gpu(self):
1477  self._test_new_device((), 1)
1478  self._test_new_device((30, 20), 1)
1479  self._test_new_device((30, 20, 10), 1)
1480  self._test_new_device((30, 20, 10, 0), 1)
1481 
1482  def test_new(self):
1483  def test_shape(sparse_dims, nnz, with_size):
1484  x, indices, values = self._gen_sparse(sparse_dims, nnz, with_size)
1485  if not x.is_cuda:
1486  # CUDA sparse tensors currently requires the size to be
1487  # specified if nDimV > 0
1488  self.assertEqual(x.new(indices, values), x)
1489  self.assertEqual(x.new(indices, values, x.size()), x)
1490 
1491  test_shape(3, 10, 100)
1492  test_shape(3, 0, [100, 100, 0])
1493 
1494  @cpu_only # not really, but we only really want to run this once
1495  def test_factory(self):
1496  for test_empty_tensor in [True, False]:
1497  if test_empty_tensor:
1498  default_size = torch.Size([1, 3, 0])
1499  size = torch.Size([3, 3, 0])
1500  else:
1501  default_size = torch.Size([1, 3])
1502  size = torch.Size([3, 3])
1503  for include_size in [True, False]:
1504  for use_tensor_idx in [True, False]:
1505  for use_tensor_val in [True, False]:
1506  for use_cuda in ([False] if not torch.cuda.is_available() else [True, False]):
1507  # have to include size with cuda sparse tensors
1508  include_size = include_size or use_cuda
1509  dtype = torch.float64
1510  long_dtype = torch.int64
1511  device = torch.device('cpu') if not use_cuda else \
1512  torch.device(torch.cuda.device_count() - 1)
1513  indices = torch.tensor(([0], [2]), dtype=long_dtype) if use_tensor_idx else ([0], [2])
1514  if test_empty_tensor:
1515  values = self.value_empty(1, 0)
1516  else:
1517  if use_tensor_val:
1518  values = torch.tensor([1.], dtype=dtype)
1519  else:
1520  values = 1.
1521  if include_size:
1522  sparse_tensor = torch.sparse_coo_tensor(indices, values, size, dtype=dtype,
1523  device=device, requires_grad=True)
1524  else:
1525  sparse_tensor = torch.sparse_coo_tensor(indices, values, dtype=dtype,
1526  device=device, requires_grad=True)
1527  self.assertEqual(indices, sparse_tensor._indices())
1528  self.assertEqual(values, sparse_tensor._values())
1529  self.assertEqual(size if include_size else default_size, sparse_tensor.size())
1530  self.assertEqual(dtype, sparse_tensor.dtype)
1531  if use_cuda:
1532  self.assertEqual(device, sparse_tensor._values().device)
1533  self.assertEqual(True, sparse_tensor.requires_grad)
1534 
1535  def test_factory_size_check(self):
1536  indices = self.index_tensor([[1, 2],
1537  [0, 2]])
1538  values = self.value_tensor([.5, .5])
1539  sizes = torch.Size([2, 3])
1540  with self.assertRaisesRegex(RuntimeError, "size is inconsistent with indices"):
1541  torch.sparse_coo_tensor(indices, values, sizes)
1542 
1543  indices.fill_(-1)
1544  with self.assertRaisesRegex(RuntimeError, "found negative index"):
1545  torch.sparse_coo_tensor(indices, values, sizes)
1546 
1547  indices = self.index_tensor([[1, 2],
1548  [0, 2]])
1549  values = self.value_empty(2, 1, 0)
1550  sizes = torch.Size([2, 3, 1, 0])
1551  with self.assertRaisesRegex(RuntimeError, "size is inconsistent with indices"):
1552  torch.sparse_coo_tensor(indices, values, sizes)
1553 
1554  indices = self.index_tensor([[1, 2],
1555  [0, 2]])
1556  values = self.value_empty(2, 2, 2)
1557  sizes = torch.Size([0, 0, 2, 2])
1558  with self.assertRaisesRegex(RuntimeError, "size is inconsistent with indices"):
1559  torch.sparse_coo_tensor(indices, values, sizes)
1560 
1561  indices = self.index_tensor([[1, 2],
1562  [0, 2]])
1563  values = self.value_tensor([[1, 1, 1], [1, 1, 1]])
1564  sizes = torch.Size([3, 3, 2])
1565  with self.assertRaisesRegex(RuntimeError, "values has incorrect size"):
1566  torch.sparse_coo_tensor(indices, values, sizes)
1567 
1568  indices = self.index_tensor([[1, 2],
1569  [0, 2]])
1570  values = self.value_empty(2, 1, 0)
1571  sizes = torch.Size([3, 3, 2, 0])
1572  with self.assertRaisesRegex(RuntimeError, "values has incorrect size"):
1573  torch.sparse_coo_tensor(indices, values, sizes)
1574 
1575  def test_factory_default(self):
1576  tensor = self.legacy_sparse_tensor()
1577  expected_indices = self.index_tensor([[]])
1578  expected_size = torch.Size([0])
1579  self.assertEqual(tensor._indices(), expected_indices)
1580  self.assertEqual(tensor.shape, expected_size)
1581 
1582  def test_factory_empty_indices(self):
1583  device = 'cuda' if self.is_cuda else 'cpu'
1584  tensor = self.legacy_sparse_tensor()
1585  expected_indices = torch.empty((1, 0), dtype=torch.long, device=device)
1586  self.assertEqual(tensor._indices(), expected_indices)
1587 
1588  tensor = torch.sparse_coo_tensor(torch.Size([2, 0]), device=device)
1589  expected_indices = torch.empty((2, 0), dtype=torch.long, device=device)
1590  self.assertEqual(tensor._indices(), expected_indices)
1591 
1592  tensor = torch.sparse_coo_tensor(torch.Size([2, 2, 0]), device=device)
1593  expected_indices = torch.empty((3, 0), dtype=torch.long, device=device)
1594  self.assertEqual(tensor._indices(), expected_indices)
1595 
1596  tensor = torch.sparse_coo_tensor(torch.Size([2, 2, 0, 0]), device=device)
1597  expected_indices = torch.empty((4, 0), dtype=torch.long, device=device)
1598  self.assertEqual(tensor._indices(), expected_indices)
1599 
1600  def test_factory_nnz(self):
1601  indices = self.index_tensor([[0]]) # (sparse_dim, nnz): (1, 1)
1602  values = self.value_tensor([[1, 1], [1, 1]]) # (nnz, ...): (2, 2)
1603  sizes = torch.Size([2, 2])
1604  with self.assertRaisesRegex(RuntimeError, "indices and values must have same nnz"):
1605  torch.sparse_coo_tensor(indices, values, sizes)
1606 
1607  indices = self.index_tensor([[0]]) # (sparse_dim, nnz): (1, 1)
1608  values = self.value_empty(2, 0) # (nnz, ...): (2, 0)
1609  sizes = torch.Size([2, 0])
1610  with self.assertRaisesRegex(RuntimeError, "indices and values must have same nnz"):
1611  torch.sparse_coo_tensor(indices, values, sizes)
1612 
1613  def test_factory_nnz_zero(self):
1614  def test_shape(i_shape, v_shape, size, expected_size):
1615  device = 'cuda' if self.is_cuda else 'cpu'
1616  if size:
1617  t = torch.sparse_coo_tensor(torch.empty(i_shape), torch.empty(v_shape), torch.Size(size), device=device)
1618  else:
1619  t = torch.sparse_coo_tensor(torch.empty(i_shape), torch.empty(v_shape), device=device)
1620  expected_indices = torch.empty(i_shape, device=device)
1621  expected_values = torch.empty(v_shape, device=device)
1622  expected_size = torch.Size(expected_size)
1623  self.assertEqual(t._indices(), expected_indices)
1624  self.assertEqual(t._values(), expected_values)
1625  self.assertEqual(t.size(), expected_size)
1626 
1627  test_shape([1, 0], [0, 2, 4, 0], None, [0, 2, 4, 0])
1628  test_shape([3, 0], [0, 2, 4, 0], None, [0, 0, 0, 2, 4, 0])
1629  test_shape([1, 0], [0, 2, 4, 0], [0, 2, 4, 0], [0, 2, 4, 0])
1630  test_shape([3, 0], [0, 2, 4, 0], [0, 0, 0, 2, 4, 0], [0, 0, 0, 2, 4, 0])
1631  test_shape([3, 0], [0, 2, 4, 0], [1, 2, 3, 2, 4, 0], [1, 2, 3, 2, 4, 0])
1632 
1633  def test_factory_dense_dim(self):
1634  indices = self.index_tensor([[0]])
1635  values = self.value_tensor([[[1, 1, 1], [1, 1, 1]]])
1636  sizes = torch.Size([1, 3, 4])
1637  with self.assertRaisesRegex(RuntimeError, "values has incorrect size"):
1638  torch.sparse_coo_tensor(indices, values, sizes)
1639 
1640  indices = self.index_tensor([[0]])
1641  values = self.value_empty(1, 2, 3, 0)
1642  sizes = torch.Size([1, 3, 4, 0])
1643  with self.assertRaisesRegex(RuntimeError, "values has incorrect size"):
1644  torch.sparse_coo_tensor(indices, values, sizes)
1645 
1646  @cpu_only
1647  def test_factory_type_inference(self):
1648  t = torch.sparse_coo_tensor(torch.tensor(([0], [2])), torch.tensor([1.], dtype=torch.float32))
1649  self.assertEqual(torch.float32, t.dtype)
1650  t = torch.sparse_coo_tensor(torch.tensor(([0], [2])), torch.tensor([1.], dtype=torch.float64))
1651  self.assertEqual(torch.float64, t.dtype)
1652  t = torch.sparse_coo_tensor(torch.tensor(([0], [2])), torch.tensor([1]))
1653  self.assertEqual(torch.int64, t.dtype)
1654 
1655  t = torch.sparse_coo_tensor(torch.tensor(([0], [2])), torch.FloatTensor(1, 0))
1656  self.assertEqual(torch.float32, t.dtype)
1657  t = torch.sparse_coo_tensor(torch.tensor(([0], [2])), torch.DoubleTensor(1, 0))
1658  self.assertEqual(torch.float64, t.dtype)
1659  t = torch.sparse_coo_tensor(torch.tensor(([0], [2])), torch.LongTensor(1, 0))
1660  self.assertEqual(torch.int64, t.dtype)
1661 
1662  @cuda_only
1663  def test_factory_device_type_inference(self):
1664  # both indices/values are CUDA
1665  shape = (1, 3)
1666  for indices_device in ['cuda', 'cpu']:
1667  for values_device in ['cuda', 'cpu']:
1668  for sparse_device in ['cuda', 'cpu', None]:
1669  for test_empty_tensor in [True, False]:
1670  if test_empty_tensor:
1671  t = torch.sparse_coo_tensor(torch.tensor(([0], [2]), device=indices_device),
1672  self.value_empty(1, 0).to(values_device),
1673  (1, 3, 0), device=sparse_device)
1674  else:
1675  t = torch.sparse_coo_tensor(torch.tensor(([0], [2]), device=indices_device),
1676  torch.tensor([1.], device=values_device),
1677  (1, 3), device=sparse_device)
1678  should_be_cuda = sparse_device == 'cuda' or (sparse_device is None and values_device == 'cuda')
1679  self.assertEqual(should_be_cuda, t.is_cuda)
1680 
1681  @cpu_only
1682  def test_factory_copy(self):
1683  def test_tensor(indices, values, indices_equal, values_equal):
1684  sparse_tensor = torch.sparse_coo_tensor(indices, values, dtype=torch.float64)
1685  if indices_equal:
1686  self.assertEqual(indices.data_ptr(), sparse_tensor._indices().data_ptr())
1687  else:
1688  self.assertNotEqual(indices.data_ptr(), sparse_tensor._indices().data_ptr())
1689  if values_equal:
1690  self.assertEqual(values.data_ptr(), sparse_tensor._values().data_ptr())
1691  else:
1692  self.assertNotEqual(values.data_ptr(), sparse_tensor._values().data_ptr())
1693 
1694  # both correct
1695  indices = torch.tensor(([0], [2]), dtype=torch.int64)
1696  values = torch.tensor([1.], dtype=torch.float64)
1697  test_tensor(indices, values, True, True)
1698 
1699  indices = torch.tensor(([0], [2]), dtype=torch.int64)
1700  values = torch.DoubleTensor(1, 0)
1701  test_tensor(indices, values, True, True)
1702 
1703  # only indices correct
1704  indices = torch.tensor(([0], [2]), dtype=torch.int64)
1705  values = torch.tensor([1.], dtype=torch.float32)
1706  test_tensor(indices, values, True, False)
1707 
1708  indices = torch.tensor(([0], [2]), dtype=torch.int64)
1709  values = torch.FloatTensor(1, 0)
1710  test_tensor(indices, values, True, True) # An empty tensor's data_ptr is always equal to 0
1711 
1712  # only values correct
1713  indices = torch.tensor(([0], [2]), dtype=torch.int32)
1714  values = torch.tensor([1.], dtype=torch.float64)
1715  test_tensor(indices, values, False, True)
1716 
1717  indices = torch.tensor(([0], [2]), dtype=torch.int32)
1718  values = torch.DoubleTensor(1, 0)
1719  test_tensor(indices, values, False, True)
1720 
1721  # neither correct
1722  indices = torch.tensor(([0], [2]), dtype=torch.int32)
1723  values = torch.tensor([1.], dtype=torch.float32)
1724  test_tensor(indices, values, False, False)
1725 
1726  indices = torch.tensor(([0], [2]), dtype=torch.int32)
1727  values = torch.FloatTensor(1, 0)
1728  test_tensor(indices, values, False, True) # An empty tensor's data_ptr is always equal to 0
1729 
1730  @cpu_only # just run once, we test both cpu and cuda
1731  def test_constructor_device_legacy(self):
1732  i = torch.tensor([[0, 1, 1], [2, 0, 2]])
1733  v = torch.tensor([3., 4., 5.])
1734  size = torch.Size([2, 3])
1735 
1736  self.assertRaises(RuntimeError, lambda: torch.sparse.FloatTensor(device='cuda'))
1737  self.assertRaises(RuntimeError, lambda: torch.sparse.FloatTensor(i, v, device='cuda'))
1738  self.assertRaises(RuntimeError, lambda: torch.sparse.FloatTensor(i, v, size, device='cuda'))
1739  self.assertRaises(RuntimeError, lambda: torch.sparse.FloatTensor(torch.Size([2, 3, 4]), device='cuda'))
1740 
1741  x = torch.sparse_coo_tensor(i, v, size, device='cpu')
1742  self.assertRaises(RuntimeError, lambda: x.new(device='cuda'))
1743  self.assertRaises(RuntimeError, lambda: x.new(i, v, device='cuda'))
1744  self.assertRaises(RuntimeError, lambda: x.new(i, v, size, device='cuda'))
1745  self.assertRaises(RuntimeError, lambda: x.new(torch.Size([2, 3, 4]), device='cuda'))
1746 
1748  self.assertRaises(RuntimeError, lambda: torch.cuda.sparse.FloatTensor(device='cpu'))
1749  self.assertRaises(RuntimeError, lambda: torch.cuda.sparse.FloatTensor(i, v, device='cpu'))
1750  self.assertRaises(RuntimeError, lambda: torch.cuda.sparse.FloatTensor(i, v, size, device='cpu'))
1751  self.assertRaises(RuntimeError, lambda: torch.cuda.sparse.FloatTensor(torch.Size([2, 3, 4]), device='cpu'))
1752 
1753  x = torch.sparse_coo_tensor(i, v, size, device='cuda')
1754  self.assertRaises(RuntimeError, lambda: x.new(device='cpu'))
1755  self.assertRaises(RuntimeError, lambda: x.new(i, v, device='cpu'))
1756  self.assertRaises(RuntimeError, lambda: x.new(i, v, size, device='cpu'))
1757  self.assertRaises(RuntimeError, lambda: x.new(torch.Size([2, 3, 4]), device='cpu'))
1758 
1759  @cpu_only # not really, but we only really want to run this once
1760  def test_dtypes(self):
1761  all_sparse_dtypes = [dtype for dtype in torch.testing.get_all_dtypes() if dtype != torch.float16]
1762  do_test_dtypes(self, all_sparse_dtypes, torch.sparse_coo, torch.device('cpu'))
1764  do_test_dtypes(self, all_sparse_dtypes, torch.sparse_coo, torch.device('cuda:0'))
1765 
1766  @cpu_only # not really, but we only really want to run this once
1767  def test_empty_full(self):
1768  all_sparse_dtypes = [dtype for dtype in torch.testing.get_all_dtypes() if dtype != torch.float16]
1769  do_test_empty_full(self, all_sparse_dtypes, torch.sparse_coo, torch.device('cpu'))
1770  if torch.cuda.device_count() > 0:
1771  do_test_empty_full(self, all_sparse_dtypes, torch.sparse_coo, None)
1772  do_test_empty_full(self, all_sparse_dtypes, torch.sparse_coo, torch.device('cuda:0'))
1773 
1774  def test_is_sparse(self):
1775  x = torch.randn(3, 3)
1776  self.assertFalse(x.is_sparse)
1777 
1778  x = torch.randn(3, 3, 0)
1779  self.assertFalse(x.is_sparse)
1780 
1781  x = self.legacy_sparse_tensor()
1782  self.assertTrue(x.is_sparse)
1783 
1784  x = self.sparse_empty(1, 0)
1785  self.assertTrue(x.is_sparse)
1786 
1787  def test_resize_as(self):
1788  def do_test(t):
1789  y = t.new().resize_as_(t).zero_()
1790  self.assertEqual(y.shape, t.shape)
1791  # Check that y can be added to t. Currently, this requires that
1792  # sparse_dim and dense_dim match.
1793  self.assertEqual(t, t + y)
1794 
1795  do_test(self.legacy_sparse_tensor())
1796  do_test(self.sparse_empty(3, 0))
1797  do_test(self.sparse_empty(3, 3))
1798 
1799  def _test_resize_shape(self, x_i, x_v, x_size, y_i, y_v, y_size):
1800  x_v_numel = torch.zeros(x_v).numel()
1801  y_v_numel = torch.zeros(y_v).numel()
1802  x = torch.sparse_coo_tensor(torch.zeros(x_i),
1803  torch.arange(x_v_numel).resize_(x_v).to(torch.float),
1804  torch.Size(x_size))
1805  x_dense = x.to_dense()
1806  y = torch.sparse_coo_tensor(torch.zeros(y_i),
1807  torch.ones(y_v).to(torch.float),
1808  torch.Size(y_size))
1809  y_dense = y.to_dense()
1810  x.resize_as_(y)
1811  x_dense.resize_as_(y_dense)
1812  self.assertEqual(x.shape, y.shape)
1813  self.assertEqual(x.sparse_dim(), y.sparse_dim())
1814  self.assertEqual(x.dense_dim(), y.dense_dim())
1815  self.assertEqual(x.shape, x_dense.shape)
1816  self.assertEqual(y.shape, y_dense.shape)
1817  # Here we make sure that the original data are preserved after resizing
1818  self.assertEqual(x.to_dense().view(-1)[0:x_v_numel].view(x_v),
1819  x_dense.view(-1)[0:x_v_numel].view(x_v))
1820 
1821  def test_resize(self):
1822  # 1. Expand the size of some dense dimensions [Supported]
1823  self._test_resize_shape([1, 1], [1, 2, 3], [2, 2, 3],
1824  [1, 1], [1, 2, 4], [2, 2, 4])
1825 
1826  self._test_resize_shape([1, 1], [1, 2, 0], [2, 2, 0],
1827  [1, 1], [1, 2, 4], [2, 2, 4])
1828 
1829  # 2. Expand the size of some sparse dimensions [Supported]
1830  self._test_resize_shape([1, 1], [1, 2, 3], [2, 2, 3],
1831  [1, 1], [1, 2, 3], [4, 2, 3])
1832 
1833  # 3. Change the shapes of both sparse and dense dimensions when nnz is zero [Supported]
1834  self._test_resize_shape([1, 0], [0, 2, 3], [2, 2, 3],
1835  [2, 0], [0, 2, 4, 5], [1, 1, 2, 4, 5])
1836 
1837  self._test_resize_shape([1, 0], [0, 2, 3], [2, 2, 3],
1838  [2, 0], [0, 2, 4, 0], [1, 1, 2, 4, 0])
1839 
1840  # 4. Add dims to dense dimensions [Not Supported]
1841  with self.assertRaisesRegex(RuntimeError, "changing the number of dense dimensions"):
1842  self._test_resize_shape([1, 1], [1, 2, 3], [2, 2, 3],
1843  [1, 1], [1, 2, 3, 4], [2, 2, 3, 4])
1844 
1845  with self.assertRaisesRegex(RuntimeError, "changing the number of dense dimensions"):
1846  self._test_resize_shape([1, 1], [1, 2, 3], [2, 2, 3],
1847  [1, 1], [1, 2, 3, 0], [2, 2, 3, 0])
1848 
1849  # 5. Remove dims from dense dimensions [Not Supported]
1850  with self.assertRaisesRegex(RuntimeError, "changing the number of dense dimensions"):
1851  self._test_resize_shape([1, 1], [1, 2, 3], [2, 2, 3],
1852  [1, 1], [1, 2], [2, 2])
1853 
1854  # 6. Change the number of sparse dimensions on a non-empty sparse tensor [Not Supported]
1855  with self.assertRaisesRegex(RuntimeError, "changing the number of sparse dimensions"):
1856  self._test_resize_shape([1, 1], [1, 2, 3], [2, 2, 3],
1857  [2, 1], [1, 2, 3], [1, 2, 2, 3])
1858 
1859  # 7. Shrink the size of some sparse dimensions on a non-empty sparse tensor [Not Supported]
1860  with self.assertRaisesRegex(RuntimeError, "shrinking the size of sparse dimensions"):
1861  self._test_resize_shape([1, 1], [1, 2, 3], [2, 2, 3],
1862  [1, 1], [1, 2, 3], [1, 2, 3])
1863 
1864  # 8. Shrink the size of some dense dimensions on a non-empty sparse tensor [Not Supported]
1865  with self.assertRaisesRegex(RuntimeError, "shrinking the size of dense dimensions"):
1866  self._test_resize_shape([1, 1], [1, 2, 3], [2, 2, 3],
1867  [1, 1], [1, 2, 2], [2, 2, 2])
1868 
1869  with self.assertRaisesRegex(RuntimeError, "shrinking the size of dense dimensions"):
1870  self._test_resize_shape([1, 1], [1, 2, 3], [2, 2, 3],
1871  [1, 1], [1, 2, 0], [2, 2, 0])
1872 
1873  def test_is_nonzero(self):
1874  self.assertTrue(torch.sparse_coo_tensor(([0],), 1., (1,)).is_nonzero())
1875  self.assertFalse(torch.sparse_coo_tensor(([0],), 0., (1,)).is_nonzero())
1876  self.assertFalse(torch.sparse_coo_tensor(([0], [0]), 0., (1, 1)).is_nonzero())
1877  self.assertFalse(torch.sparse_coo_tensor(([0, 0],), (0., 0.), (1,)).is_nonzero())
1878  self.assertFalse(torch.sparse_coo_tensor(([0, 0],), (-1., 1.), (1,)).is_nonzero())
1879  self.assertTrue(torch.sparse_coo_tensor(torch.zeros(0, 1), 12.3, []).is_nonzero()) # scalar sparse tensor
1880  with self.assertRaisesRegex(RuntimeError, "bool value of Tensor with no values is ambiguous"):
1881  torch.sparse_coo_tensor(([0, 1],), self.value_empty(2, 0), (4, 0)).is_nonzero()
1882 
1883  def test_allow_tensor_metadata_change(self):
1884  def do_test(t):
1885  with self.assertRaisesRegex(
1886  RuntimeError,
1887  "raw_resize_ is not allowed on Tensor created from .data or .detach()"):
1888  t.transpose_(0, 1)
1889  with self.assertRaisesRegex(
1890  RuntimeError,
1891  "resize_ is not allowed on Tensor created from .data or .detach()"):
1892  t.resize_as_(self.sparse_empty(3, 3))
1893  with self.assertRaisesRegex(
1894  RuntimeError,
1895  "resize_and_clear_ is not allowed on Tensor created from .data or .detach()"):
1896  t.mul_(t)
1897  with self.assertRaisesRegex(
1898  RuntimeError,
1899  "set_coalesced is not allowed on Tensor created from .data or .detach()"):
1900  t._coalesced_(True)
1901  with self.assertRaisesRegex(
1902  RuntimeError,
1903  "set_indices_and_values_unsafe is not allowed on Tensor created from .data or .detach()"):
1904  a = self.sparse_tensor(torch.tensor([[0, 1, 1], [2, 0, 2]]), torch.tensor([3., 4., 5.])).data
1905  a.add_(a)
1906  with self.assertRaisesRegex(
1907  RuntimeError,
1908  "resize_and_clear_ is not allowed on Tensor created from .data or .detach()"):
1909  a.zero_()
1910  with self.assertRaisesRegex(
1911  RuntimeError,
1912  "resize_ is not allowed on Tensor created from .data or .detach()"):
1913  a.copy_(self.sparse_empty(3, 3))
1914 
1915  do_test(self.sparse_empty(3, 0).data)
1916  do_test(self.sparse_empty(3, 0).detach())
1917 
1918 
1920  def setUp(self):
1921  super(TestUncoalescedSparse, self).setUp()
1922  self.is_uncoalesced = True
1923 
1924 
1925 @unittest.skipIf(not TEST_CUDA, 'CUDA not available')
1927  def setUp(self):
1928  super(TestCudaSparse, self).setUp()
1929  self.is_cuda = True
1930  self.device = 'cuda'
1931  self.legacy_sparse_tensor = torch.cuda.sparse.DoubleTensor
1932 
1933 
1934 @unittest.skipIf(not TEST_CUDA, 'CUDA not available')
1936  def setUp(self):
1937  super(TestCudaUncoalescedSparse, self).setUp()
1938  self.is_uncoalesced = True
1939 
1940 
1942  @unittest.skipIf(not TEST_CUDA, 'CUDA not available')
1943  def test_cuda_from_cpu(self):
1944  with self.assertRaisesRegex(
1945  RuntimeError,
1946  "backend of indices \\(CUDA\\) must match backend of values \\(CPU\\)"):
1947  torch.sparse.FloatTensor(torch.zeros(1, 4).long().cuda(),
1948  torch.randn(4, 4, 4),
1949  [3, 4, 4])
1950 
1951  with self.assertRaisesRegex(
1952  RuntimeError,
1953  "backend of indices \\(CUDA\\) must match backend of values \\(CPU\\)"):
1954  torch.sparse.FloatTensor(torch.zeros(1, 4).long().cuda(),
1955  torch.randn(4, 4, 4, 0),
1956  [3, 4, 4, 0])
1957 
1958  with self.assertRaisesRegex(
1959  RuntimeError,
1960  "backend of indices \\(CUDA\\) must match backend of values \\(CPU\\)"):
1961  torch.sparse.FloatTensor(torch.LongTensor(1, 0).cuda(),
1962  torch.randn(0, 4, 4, 0),
1963  [0, 4, 4, 0])
1964 
1965  @unittest.skipIf(not TEST_CUDA, 'CUDA not available')
1966  def test_cuda_sparse_cpu_dense_add(self):
1967  x = torch.zeros(3, 4, 4)
1968  sparse_y = torch.cuda.sparse.FloatTensor(torch.zeros(1, 4).long().cuda(),
1969  torch.randn(4, 4, 4).cuda(),
1970  [3, 4, 4])
1971  with self.assertRaisesRegex(RuntimeError, "add: expected 'other' to be a CPU tensor\\, but got a CUDA tensor"):
1972  x + sparse_y
1973 
1974  x = torch.zeros(3, 4, 4, 0)
1975  sparse_y = torch.cuda.sparse.FloatTensor(torch.zeros(1, 4).long().cuda(),
1976  torch.randn(4, 4, 4, 0).cuda(),
1977  [3, 4, 4, 0])
1978  with self.assertRaisesRegex(RuntimeError, "add: expected 'other' to be a CPU tensor\\, but got a CUDA tensor"):
1979  x + sparse_y
1980 
1981  x = torch.zeros(0, 4, 4, 0)
1982  sparse_y = torch.cuda.sparse.FloatTensor(torch.LongTensor(1, 0).cuda(),
1983  torch.randn(0, 4, 4, 0).cuda(),
1984  [0, 4, 4, 0])
1985  with self.assertRaisesRegex(RuntimeError, "add: expected 'other' to be a CPU tensor\\, but got a CUDA tensor"):
1986  x + sparse_y
1987 
1988 
1989 if __name__ == '__main__':
1990  run_tests()
def assertEqual(self, x, y, prec=None, message='', allow_inf=False)
def _test_basic_ops_shape(self, nnz_x1, nnz_x2, shape_i, shape_v=None)
def _test_resize_shape(self, x_i, x_v, x_size, y_i, y_v, y_size)
def _gen_sparse(self, sparse_dim, nnz, with_size)
Definition: test_sparse.py:65
def _test_sparse_mask_fixed(self)
def addmm(mat, mat1, mat2, beta=1, alpha=1)
Definition: __init__.py:11
def assertExpected(self, s, subname=None)
def is_available()
Definition: __init__.py:45
def device_count()
Definition: __init__.py:341
def _test_zeros_like(self, nnzs, template_shape_i, template_shape_v=None)
def assertNotEqual(self, x, y, prec=None, message='')
def _all_narrow_combs(self, shape)
def get_all_dtypes()
Definition: __init__.py:89
def assert_uncoalesced(self, x)
Definition: test_sparse.py:76
def _test_sparse_mask_hybrid_fixed(self)
def randn(self, args, kwargs)
Definition: test_sparse.py:90
def _test_zeros(self, nnzs, shape, out_shape_i, out_shape_v=None)
def _test_spadd_shape(self, nnz, shape_i, shape_v=None)
Definition: test_sparse.py:916
def safeCoalesce(self, t)
def _test_narrow(self, input, narrow_args)
def _test_log1p_tensor(self, input, dense_tensor)
def genSparseTensor(self, size, sparse_dim, nnz, is_uncoalesced, device='cpu')
def safeToDense(self, t)
def sum(input, dim=None, dtype=None)
Definition: __init__.py:70
def _test_new_device(self, size, device)
def _test_sparse_mask_shape(self, nnz_x1, nnz_x2, shape_i, shape_v=None)
def mm(mat1, mat2)
Definition: __init__.py:28