mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-14 20:57:59 +00:00
Generator of tensor inputs with variable layout and structure (batch/non-batch, hybrid/non-hybrid, block/non-block) (#88914)
This PR introduces `TestCase.generate_simple_inputs` method that is an improved and generalized version of the `TestSparseCompressed._generate_small_inputs` method. Pull Request resolved: https://github.com/pytorch/pytorch/pull/88914 Approved by: https://github.com/cpuhrsch
This commit is contained in:
Pearu Peterson
PyTorch MergeBot
parent
275ade6371
commit
90bed8874f
11 changed files with 17816 additions and 30112 deletions
File diff suppressed because it is too large
Load diff
File diff suppressed because it is too large
Load diff
File diff suppressed because it is too large
Load diff
File diff suppressed because it is too large
Load diff
File diff suppressed because it is too large
Load diff
File diff suppressed because it is too large
Load diff
File diff suppressed because it is too large
Load diff
File diff suppressed because it is too large
Load diff
|
|
@ -4049,6 +4049,52 @@ class TestSparseMeta(TestCase):
|
|||
self.assertEqual(r.values(), torch.empty(0, 4, device='meta'))
|
||||
|
||||
|
||||
class TestSparseAny(TestCase):
|
||||
|
||||
def test_generate_simple_inputs(self):
|
||||
# Temporarily disable BSC and BSC layouts as these don't support select yet, see the next PR in the stack.
|
||||
layouts = [torch.strided, torch.sparse_coo, torch.sparse_csr, torch.sparse_csc, torch.sparse_bsr, torch.sparse_bsc][:-2]
|
||||
|
||||
tested_combinations = set()
|
||||
for tensors in zip(*map(self.generate_simple_inputs, layouts)):
|
||||
for i, t in enumerate(tensors):
|
||||
self.assertEqual(t.layout, layouts[i])
|
||||
|
||||
# all layouts must produce semantically the same tensors
|
||||
self.assertEqual(t, tensors[0])
|
||||
|
||||
if t.layout is torch.strided:
|
||||
is_hybrid = None
|
||||
else:
|
||||
is_hybrid = t.dense_dim() > 0
|
||||
if t.layout in {torch.sparse_csr, torch.sparse_bsr}:
|
||||
is_batch = t.crow_indices().ndim > 1
|
||||
elif t.layout in {torch.sparse_csc, torch.sparse_bsc}:
|
||||
is_batch = t.ccol_indices().ndim > 1
|
||||
else:
|
||||
is_batch = None
|
||||
if t.layout in {torch.sparse_bsr, torch.sparse_bsc}:
|
||||
blocksize = t.values().shape[1:3]
|
||||
nontrivial_blocksize = 1 not in blocksize
|
||||
else:
|
||||
nontrivial_blocksize = None
|
||||
tested_combinations.add((t.layout, is_hybrid, is_batch, nontrivial_blocksize))
|
||||
|
||||
# Ensure that the inputs generation covers all layout,
|
||||
# non-hybrid/hybrid, and non-batch/batch combinations:
|
||||
for layout in layouts:
|
||||
for is_hybrid in [False, True]:
|
||||
if layout is torch.strided:
|
||||
is_hybrid = None
|
||||
for is_batch in [False, True]:
|
||||
if layout in {torch.sparse_coo, torch.strided}:
|
||||
is_batch = None
|
||||
for nontrivial_blocksize in [False, True]:
|
||||
if layout not in {torch.sparse_bsr, torch.sparse_bsc}:
|
||||
nontrivial_blocksize = None
|
||||
key = (layout, is_hybrid, is_batch, nontrivial_blocksize)
|
||||
assert key in tested_combinations, key
|
||||
|
||||
|
||||
# e.g., TestSparseUnaryUfuncsCPU and TestSparseUnaryUfuncsCUDA
|
||||
instantiate_device_type_tests(TestSparseUnaryUfuncs, globals(), except_for='meta')
|
||||
|
|
@ -4058,5 +4104,7 @@ instantiate_device_type_tests(TestSparseMaskedReductions, globals(), except_for=
|
|||
# e.g., TestSparseCPU and TestSparseCUDA
|
||||
instantiate_device_type_tests(TestSparse, globals(), except_for='meta')
|
||||
|
||||
instantiate_device_type_tests(TestSparseAny, globals(), except_for='meta')
|
||||
|
||||
if __name__ == '__main__':
|
||||
run_tests()
|
||||
|
|
|
|||
|
|
@ -1,6 +1,5 @@
|
|||
# Owner(s): ["module: sparse"]
|
||||
|
||||
import copy
|
||||
import torch
|
||||
import random
|
||||
import itertools
|
||||
|
|
@ -198,147 +197,6 @@ class TestSparseCompressed(TestCase):
|
|||
device = self.device_type
|
||||
return self.genSparseCompressedTensor(size, nnz, device=device, dtype=dtype, index_dtype=index_dtype, layout=layout)
|
||||
|
||||
def _generate_small_inputs_utils(self, layout, device=None, dtype=None):
|
||||
|
||||
def shape(shape, basedim=0, blocksize=(1, 1), dense_shape=()):
|
||||
# Below, we define compressed and plain indices that
|
||||
# correspond to row compressed tensors. In order to reuse
|
||||
# the indices tensors for column compressed tensors, we
|
||||
# swap the row and columns in shape dims (basedim and
|
||||
# basedim + 1, respectively) to obtain the correct shape
|
||||
# for column compressed tensors. Batch and dense
|
||||
# dimensions remain as they are.
|
||||
#
|
||||
# Similarly, we reuse indices of non-block tensors for
|
||||
# block tensors, that means, we'll need to multiply the
|
||||
# base shape of the non-block tensor with blocksize to get
|
||||
# the base shape of a block tensor.
|
||||
if layout is torch.sparse_csc:
|
||||
shape = shape[:basedim] + (shape[basedim + 1], shape[basedim]) + shape[basedim + 2:]
|
||||
elif layout is torch.sparse_bsc:
|
||||
shape = shape[:basedim] + (shape[basedim + 1] * blocksize[1], shape[basedim] * blocksize[0]) + shape[basedim + 2:]
|
||||
elif layout is torch.sparse_bsr:
|
||||
shape = shape[:basedim] + (shape[basedim] * blocksize[0], shape[basedim + 1] * blocksize[1]) + shape[basedim + 2:]
|
||||
return shape
|
||||
|
||||
def values(lst, basedim=0, blocksize=(1, 1), densesize=(), device=device, dtype=dtype):
|
||||
# Below, we define values for non-blocked and non-hybrid
|
||||
# tensors. To reuse these for blocked tensors, we replace
|
||||
# all values in lst with a double-list that "shape"
|
||||
# corresponds to blocksize.
|
||||
# To support hybrid tensors, the values in lst are further
|
||||
# replaced with a N-list where N==len(densesize) and the
|
||||
# shape corresponds to densesize.
|
||||
|
||||
max_val = torch.iinfo(dtype).max if dtype in [torch.int16, torch.int8, torch.uint8] else None
|
||||
|
||||
def list_add(lst, value):
|
||||
# recursively add a value to lst items
|
||||
if isinstance(lst, list):
|
||||
return [list_add(item, value) for item in lst]
|
||||
rc = lst + value
|
||||
return rc if max_val is None else (rc % max_val)
|
||||
|
||||
def stretch_values(value, bdim, values_item_shape):
|
||||
# replace a value with a new value that extends the
|
||||
# dimensionality of the value by
|
||||
# len(values_item_shape) from right. The left
|
||||
# dimensions up to bdim are considered as batch
|
||||
# dimensions.
|
||||
if not values_item_shape:
|
||||
return value
|
||||
if isinstance(value, list) and bdim >= 0:
|
||||
return [stretch_values(item, bdim - 1, values_item_shape) for item in value]
|
||||
new_value = functools.reduce(lambda x, dims: [copy.deepcopy(x) for _ in range(dims)],
|
||||
reversed(values_item_shape), None)
|
||||
for p in itertools.product(*map(list, map(range, values_item_shape))):
|
||||
row = functools.reduce(lambda x, i: x.__getitem__(i), p[:-1], new_value)
|
||||
row[p[-1]] = list_add(value, sum([i * 10 ** d for d, i in enumerate(p)]))
|
||||
return new_value
|
||||
|
||||
if layout is torch.sparse_bsr:
|
||||
values_item_shape = blocksize + densesize
|
||||
elif layout is torch.sparse_bsc:
|
||||
values_item_shape = tuple(reversed(blocksize)) + densesize
|
||||
else:
|
||||
values_item_shape = densesize
|
||||
|
||||
if not lst:
|
||||
return torch.tensor(lst, device=device, dtype=dtype).reshape(0, *values_item_shape)
|
||||
|
||||
lst = stretch_values(lst, basedim, values_item_shape)
|
||||
|
||||
return torch.tensor(lst, device=device, dtype=dtype)
|
||||
|
||||
return shape, values
|
||||
|
||||
def _generate_small_inputs(self, layout, device=None, dtype=None, index_dtype=None,
|
||||
enable_batched=True, enable_hybrid=True):
|
||||
"""Generator of inputs to sparse compressed tensor factory functions.
|
||||
|
||||
The input is defined as a 4-tuple:
|
||||
compressed_indices, plain_indices, values, expected_size_from_shape_inference
|
||||
"""
|
||||
if index_dtype is None:
|
||||
index_dtype = torch.int64
|
||||
|
||||
shape, values = self._generate_small_inputs_utils(layout, device, dtype)
|
||||
|
||||
# a regular tensor
|
||||
yield (torch.tensor([0, 2, 4], device=device, dtype=index_dtype),
|
||||
torch.tensor([0, 1, 0, 2], device=device, dtype=index_dtype),
|
||||
values([1, 2, 3, 4], 0, (2, 1)),
|
||||
shape((2, 3), 0, (2, 1)))
|
||||
|
||||
# a tensor with zero dimensions
|
||||
yield (torch.tensor([0, ], device=device, dtype=index_dtype),
|
||||
torch.tensor([], device=device, dtype=index_dtype),
|
||||
values([], 0, (2, 1)),
|
||||
shape((0, 0), 0, (2, 1)))
|
||||
|
||||
if enable_batched:
|
||||
# a batched tensor with one batch dimension
|
||||
yield (torch.tensor([[0, 2, 4], [0, 3, 4]], device=device, dtype=index_dtype),
|
||||
torch.tensor([[0, 1, 0, 1], [0, 1, 2, 0]], device=device, dtype=index_dtype),
|
||||
values([[1, 2, 3, 4], [5, 6, 7, 8]], 1, (1, 2)),
|
||||
shape((2, 2, 3), 1, (1, 2)))
|
||||
|
||||
# a batched tensor with two batch dimensions
|
||||
yield (torch.tensor([[[0, 2, 4], [0, 3, 4], [0, 1, 4]],
|
||||
[[0, 1, 4], [0, 2, 4], [0, 3, 4]]],
|
||||
device=device, dtype=index_dtype),
|
||||
torch.tensor([[[0, 1, 0, 1], [0, 1, 2, 0], [0, 0, 1, 2]],
|
||||
[[1, 0, 1, 2], [0, 2, 0, 1], [0, 1, 2, 1]]],
|
||||
device=device, dtype=index_dtype),
|
||||
values([[[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]],
|
||||
[[13, 14, 15, 16], [17, 18, 19, 20], [21, 22, 23, 24]]], 2, (2, 3)),
|
||||
shape((2, 3, 2, 3), 2, (2, 3)))
|
||||
|
||||
if enable_hybrid:
|
||||
# a tensor with one dense dimension
|
||||
yield (torch.tensor([0, 2, 4], device=device, dtype=index_dtype),
|
||||
torch.tensor([0, 1, 0, 2], device=device, dtype=index_dtype),
|
||||
values([1, 2, 3, 4], 0, (3, 2), (2,)),
|
||||
shape((2, 3, 2), 0, (3, 2)))
|
||||
|
||||
# a tensor with two dense dimensions
|
||||
yield (torch.tensor([0, 2, 4], device=device, dtype=index_dtype),
|
||||
torch.tensor([0, 1, 0, 2], device=device, dtype=index_dtype),
|
||||
values([1, 2, 3, 4], 0, (2, 3), (4, 2)),
|
||||
shape((2, 3, 4, 2), 0, (2, 3)))
|
||||
|
||||
if enable_batched and enable_hybrid:
|
||||
# a batched tensor with two batch dimensions and two dense dimensions
|
||||
yield (torch.tensor([[[0, 2, 4], [0, 3, 4], [0, 1, 4]],
|
||||
[[0, 1, 4], [0, 2, 4], [0, 3, 4]]],
|
||||
device=device, dtype=index_dtype),
|
||||
torch.tensor([[[0, 1, 0, 1], [0, 1, 2, 0], [0, 0, 1, 2]],
|
||||
[[1, 0, 1, 2], [0, 2, 0, 1], [0, 1, 2, 1]]],
|
||||
device=device, dtype=index_dtype),
|
||||
values([[[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]],
|
||||
[[13, 14, 15, 16], [17, 18, 19, 20], [21, 22, 23, 24]]], 2, (3, 2), (2, 1)),
|
||||
shape((2, 3, 2, 3, 2, 1), 2, (3, 2)))
|
||||
|
||||
@all_sparse_compressed_layouts()
|
||||
@onlyCPU
|
||||
def test_layout(self, layout):
|
||||
|
|
@ -352,11 +210,14 @@ class TestSparseCompressed(TestCase):
|
|||
@dtypes(*all_types_and_complex_and(torch.half, torch.bool, torch.bfloat16))
|
||||
def test_sparse_compressed_constructor(self, layout, device, dtype,
|
||||
use_factory_function, shape_and_device_inference, input_kind):
|
||||
if input_kind == 'list' and shape_and_device_inference and torch.device(device).type == 'cuda':
|
||||
# list inputs to factory/constructor function without
|
||||
# specifying device will result a sparse compressed tensor
|
||||
# on CPU. So, skip testing against cuda device as unused.
|
||||
self.skipTest("nothing to test")
|
||||
if input_kind == 'list' and shape_and_device_inference:
|
||||
if torch.device(device).type == 'cuda':
|
||||
# list inputs to factory/constructor function without
|
||||
# specifying device will result a sparse compressed tensor
|
||||
# on CPU. So, skip testing against cuda device as unused.
|
||||
self.skipTest("nothing to test")
|
||||
if dtype not in {torch.float32, torch.complex64, torch.int64, torch.bool}:
|
||||
self.skipTest("dtype not supported with list values")
|
||||
|
||||
expected_devices = [torch.device(device)]
|
||||
if TEST_CUDA and torch.device(device).type == 'cuda' and torch.cuda.device_count() >= 2 and not shape_and_device_inference:
|
||||
|
|
@ -369,29 +230,34 @@ class TestSparseCompressed(TestCase):
|
|||
torch.sparse_bsc: torch.sparse_bsc_tensor,
|
||||
}[layout]
|
||||
compressed_indices_mth, plain_indices_mth = sparse_compressed_indices_methods[layout]
|
||||
for index_dtype in [torch.int32, torch.int64]:
|
||||
if input_kind == 'list':
|
||||
index_dtypes = [torch.int64]
|
||||
else:
|
||||
index_dtypes = [torch.int32, torch.int64]
|
||||
for index_dtype in index_dtypes:
|
||||
for expected_device in expected_devices:
|
||||
for compressed_indices, plain_indices, values, size in self._generate_small_inputs(
|
||||
layout, expected_device, dtype, index_dtype):
|
||||
for (compressed_indices, plain_indices, values), kwargs in self.generate_simple_inputs(
|
||||
layout, device=expected_device, dtype=dtype, index_dtype=index_dtype,
|
||||
# skip zero-sized tensors for list inputs:
|
||||
enable_zero_sized=input_kind != 'list',
|
||||
output_tensor=False):
|
||||
size = kwargs['size']
|
||||
if shape_and_device_inference and 0 in size:
|
||||
# skip shape inference for zero-sized tensor
|
||||
# inputs because (i) the shape determined from
|
||||
# an empty list is ambiguous, and (ii) the
|
||||
# size of the plain dimension defined as
|
||||
# max(plain_indices) is undefined if
|
||||
# plain_indices has no values
|
||||
continue
|
||||
compressed_indices_expect = compressed_indices
|
||||
plain_indices_expect = plain_indices
|
||||
values_expect = values
|
||||
|
||||
if input_kind == 'list':
|
||||
if size == (0, 0):
|
||||
# for this degenerate case, plain_indices must
|
||||
# remain a tensor because
|
||||
# tensor(plain_indices) results a float dtype
|
||||
# when plain_indices is an empty list
|
||||
if index_dtype == torch.int32:
|
||||
# skip testing int32 case because
|
||||
# tensor(compressed_indices) results a
|
||||
# int64 dtype when compressed_indices is
|
||||
# [0] (a list of single int zero).
|
||||
continue
|
||||
else:
|
||||
plain_indices = plain_indices.tolist()
|
||||
compressed_indices = compressed_indices.tolist()
|
||||
plain_indices = plain_indices.tolist()
|
||||
values = values.tolist()
|
||||
if size == (0, 0) and layout in {torch.sparse_bsr, torch.sparse_bsc}:
|
||||
# in the block sparse case, values of type list needs to represent a 3-D tensor
|
||||
values = [[[]]]
|
||||
|
||||
if use_factory_function:
|
||||
if shape_and_device_inference:
|
||||
|
|
@ -407,9 +273,9 @@ class TestSparseCompressed(TestCase):
|
|||
dtype=dtype, layout=layout, device=expected_device)
|
||||
self.assertEqual(layout, sparse.layout)
|
||||
self.assertEqual(size, sparse.shape)
|
||||
self.assertEqual(compressed_indices, compressed_indices_mth(sparse))
|
||||
self.assertEqual(plain_indices, plain_indices_mth(sparse))
|
||||
self.assertEqual(values, sparse.values())
|
||||
self.assertEqual(compressed_indices_expect, compressed_indices_mth(sparse))
|
||||
self.assertEqual(plain_indices_expect, plain_indices_mth(sparse))
|
||||
self.assertEqual(values_expect, sparse.values())
|
||||
self.assertEqual(sparse.device, sparse.values().device)
|
||||
self.assertEqual(sparse.device, expected_device)
|
||||
|
||||
|
|
@ -455,10 +321,10 @@ class TestSparseCompressed(TestCase):
|
|||
@all_sparse_compressed_layouts()
|
||||
@dtypes(*all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16))
|
||||
def test_clone(self, layout, device, dtype):
|
||||
for compressed_indices, plain_indices, values, size in self._generate_small_inputs(
|
||||
layout, device, dtype, index_dtype=torch.int32):
|
||||
sparse = torch.sparse_compressed_tensor(compressed_indices, plain_indices, values, size,
|
||||
dtype=dtype, layout=layout, device=device)
|
||||
for sparse in self.generate_simple_inputs(
|
||||
layout, device=device, dtype=dtype, index_dtype=torch.int32,
|
||||
# Temporarily disable testing batch block tensors:
|
||||
enable_batch=layout in {torch.sparse_csr, torch.sparse_csc}):
|
||||
cloned_sparse = sparse.clone()
|
||||
self.assertEqual(sparse, cloned_sparse)
|
||||
|
||||
|
|
@ -467,10 +333,37 @@ class TestSparseCompressed(TestCase):
|
|||
compressed_indices_mth, plain_indices_mth = sparse_compressed_indices_methods[layout]
|
||||
printed = []
|
||||
for enable_hybrid in [False, True]:
|
||||
# using local patterns for test_print stability
|
||||
patterns = [
|
||||
# 2 x 3 batch of 3 x 2 tensors, trivial blocksize, non-hybrid/hybrid:
|
||||
([[[[1, 2, 0],
|
||||
[1, 0, 3]],
|
||||
[[1, 2, 3],
|
||||
[1, 0, 0]],
|
||||
[[1, 0, 0],
|
||||
[1, 2, 3]]],
|
||||
[[[0, 2, 0],
|
||||
[1, 2, 3]],
|
||||
[[1, 0, 3],
|
||||
[1, 2, 0]],
|
||||
[[1, 2, 3],
|
||||
[0, 2, 0]]]], [(2, 1)], [(), (4,)] if enable_hybrid else [()]),
|
||||
# tensor with non-trivial blocksize, non-hybrid/hybrid:
|
||||
([[0, 1, 0, 2, 0, 2],
|
||||
[0, 1, 0, 0, 2, 0],
|
||||
[3, 3, 3, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[0, 5, 0, 6, 6, 6],
|
||||
[5, 0, 5, 6, 6, 6],
|
||||
[0, 0, 0, 0, 8, 8],
|
||||
[7, 7, 7, 0, 8, 8]], [(2, 3)], [(), (4, 2)] if enable_hybrid else [()]),
|
||||
]
|
||||
for index_dtype in [torch.int32, torch.int64]:
|
||||
for dtype in [torch.float32, torch.float64]:
|
||||
for compressed_indices, plain_indices, values, size in self._generate_small_inputs(
|
||||
layout, device, dtype, index_dtype, enable_hybrid=enable_hybrid):
|
||||
for (compressed_indices, plain_indices, values), kwargs in self.generate_simple_inputs(
|
||||
layout, device=device, dtype=dtype, index_dtype=index_dtype, enable_hybrid=enable_hybrid,
|
||||
enable_zero_sized=False, output_tensor=False, patterns=patterns):
|
||||
size = tuple(kwargs['size'])
|
||||
block_ndim = 2 if layout in {torch.sparse_bsr, torch.sparse_bsc} else 0
|
||||
base_ndim = 2
|
||||
batch_ndim = compressed_indices.dim() - 1
|
||||
|
|
@ -647,9 +540,7 @@ class TestSparseCompressed(TestCase):
|
|||
@all_sparse_compressed_layouts('layout2')
|
||||
@dtypes(*all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16))
|
||||
def test_empty_like(self, layout, layout2, device, dtype):
|
||||
for compressed_indices, plain_indices, values, size in self._generate_small_inputs(layout):
|
||||
sparse = torch.sparse_compressed_tensor(compressed_indices, plain_indices, values, size,
|
||||
dtype=dtype, layout=layout, device=device)
|
||||
for sparse in self.generate_simple_inputs(layout):
|
||||
if layout == layout2:
|
||||
result = torch.empty_like(sparse, layout=layout2)
|
||||
compressed_indices_mth, plain_indices_mth = sparse_compressed_indices_methods[result.layout]
|
||||
|
|
@ -671,14 +562,28 @@ class TestSparseCompressed(TestCase):
|
|||
@dtypes(*all_types_and_complex_and(torch.half, torch.bool, torch.bfloat16))
|
||||
def test_validate(self, layout, device, dtype):
|
||||
for index_dtype in [torch.int32, torch.int64]:
|
||||
for compressed_indices, plain_indices, values, size in self._generate_small_inputs(
|
||||
layout, device, dtype, index_dtype, enable_batched=True, enable_hybrid=True):
|
||||
for (compressed_indices, plain_indices, values), kwargs in self.generate_simple_inputs(
|
||||
layout, device=device, dtype=dtype, index_dtype=index_dtype, output_tensor=False):
|
||||
size = kwargs['size']
|
||||
torch._validate_sparse_compressed_tensor_args(compressed_indices, plain_indices, values, size, layout)
|
||||
|
||||
def _generate_invalid_input(self, layout, device):
|
||||
from functools import partial
|
||||
|
||||
shape, values = self._generate_small_inputs_utils(layout, device=device)
|
||||
def shape(shape, basedim=0):
|
||||
blocksize = (1, 1)
|
||||
if layout is torch.sparse_csc:
|
||||
shape = shape[:basedim] + (shape[basedim + 1], shape[basedim]) + shape[basedim + 2:]
|
||||
elif layout is torch.sparse_bsc:
|
||||
shape = shape[:basedim] + (shape[basedim + 1] * blocksize[1], shape[basedim] * blocksize[0]) + shape[basedim + 2:]
|
||||
elif layout is torch.sparse_bsr:
|
||||
shape = shape[:basedim] + (shape[basedim] * blocksize[0], shape[basedim + 1] * blocksize[1]) + shape[basedim + 2:]
|
||||
return shape
|
||||
|
||||
def values(lst, device=device):
|
||||
if layout in {torch.sparse_bsr, torch.sparse_bsc}:
|
||||
lst = [[[item]] for item in lst]
|
||||
return torch.tensor(lst, device=device)
|
||||
|
||||
tensor = partial(torch.tensor, device=device)
|
||||
values = partial(values, device=device)
|
||||
|
|
@ -925,7 +830,8 @@ class TestSparseCompressed(TestCase):
|
|||
@onlyCPU
|
||||
@all_sparse_compressed_layouts()
|
||||
def test_dim(self, layout):
|
||||
for compressed_indices, plain_indices, values, size in self._generate_small_inputs(layout):
|
||||
for (compressed_indices, plain_indices, values), kwargs in self.generate_simple_inputs(layout, output_tensor=False):
|
||||
size = kwargs['size']
|
||||
batch_dim = compressed_indices.dim() - 1
|
||||
sparse_dim = 2
|
||||
block_dim = 2 if layout in {torch.sparse_bsr, torch.sparse_bsc} else 0
|
||||
|
|
@ -940,10 +846,7 @@ class TestSparseCompressed(TestCase):
|
|||
@dtypes(*all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16))
|
||||
def test_to_dtype(self, layout, device, dtype):
|
||||
# to_dense does not support hybrid inputs
|
||||
input_gen = self._generate_small_inputs(layout, device=device, enable_hybrid=False)
|
||||
for compressed_indices, plain_indices, values, size in input_gen:
|
||||
sparse = torch.sparse_compressed_tensor(compressed_indices, plain_indices, values, size,
|
||||
dtype=dtype, layout=layout, device=device)
|
||||
for sparse in self.generate_simple_inputs(layout, dtype=dtype, device=device, enable_hybrid=False):
|
||||
for to_dtype in all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16):
|
||||
sparse_to_dtype = sparse.to(to_dtype)
|
||||
dense_to_dtype = sparse.to_dense().to(to_dtype)
|
||||
|
|
@ -955,10 +858,7 @@ class TestSparseCompressed(TestCase):
|
|||
def test_pickle(self, layout, dtype, device):
|
||||
import pickle
|
||||
|
||||
input_gen = self._generate_small_inputs(layout)
|
||||
for compressed_indices, plain_indices, values, size in input_gen:
|
||||
sparse = torch.sparse_compressed_tensor(compressed_indices, plain_indices, values, size,
|
||||
dtype=dtype, device=device, layout=layout)
|
||||
for sparse in self.generate_simple_inputs(layout, device=device, dtype=dtype):
|
||||
serialized = pickle.dumps(sparse)
|
||||
sparse_loaded = pickle.loads(serialized)
|
||||
|
||||
|
|
|
|||
|
|
@ -2446,6 +2446,296 @@ class TestCase(expecttest.TestCase):
|
|||
x = x.detach().clone()._coalesced_(False)
|
||||
return x, x._indices().clone(), x._values().clone()
|
||||
|
||||
def generate_simple_inputs(self, layout,
|
||||
device=None,
|
||||
dtype=None,
|
||||
index_dtype=None,
|
||||
enable_batch=True,
|
||||
enable_hybrid=True,
|
||||
enable_zero_sized=True,
|
||||
enable_batch_variable_nse=False,
|
||||
output_tensor=True,
|
||||
patterns=None):
|
||||
"""Generator of simple inputs for tensor constructors of the given layout.
|
||||
|
||||
The generated tensor inputs have the following properties:
|
||||
|
||||
- tensor shapes are minimal but not trivial
|
||||
- tensor values are sorted sequences for COO and CSR formats, e.g. [1, 2, 3, 4]
|
||||
- the generated tensors represent the same mathematical tensor for all layouts
|
||||
- the generated tensors include regular, zero-sized, and optionally, batched or/and hybrid tensors.
|
||||
|
||||
If output_tensor is True, yield tensors with the given
|
||||
layout. Otherwise, yield inputs to the corresponding tensor
|
||||
constructors:
|
||||
|
||||
- sparse compressed input is defined as
|
||||
(compressed_indices, plain_indices, values), dict(size=expected_size_from_shape_inference, device=device, dtype=dtype)
|
||||
|
||||
- sparse COO input is defined as
|
||||
(indices, values), dict(size=expected_size_from_shape_inference, device=device, dtype=dtype)
|
||||
|
||||
- strided input is defined as
|
||||
(values,), dict(device=device, dtype=dtype)
|
||||
"""
|
||||
if index_dtype is None:
|
||||
index_dtype = torch.int64
|
||||
|
||||
is_compressed_sparse_layout = layout in {torch.sparse_csr, torch.sparse_csc, torch.sparse_bsr, torch.sparse_bsc}
|
||||
|
||||
if output_tensor:
|
||||
for args, kwargs in self.generate_simple_inputs(layout, device=device, dtype=dtype, index_dtype=index_dtype,
|
||||
enable_batch=enable_batch, enable_hybrid=enable_hybrid,
|
||||
enable_zero_sized=enable_zero_sized,
|
||||
enable_batch_variable_nse=enable_batch_variable_nse,
|
||||
output_tensor=False):
|
||||
if layout is torch.strided:
|
||||
assert len(args) == 1
|
||||
size = kwargs.pop('size', None) # to ensure that a zero-sized tensor has the desired shape
|
||||
assert size is not None
|
||||
yield args[0].reshape(size)
|
||||
elif layout is torch.sparse_coo:
|
||||
yield torch.sparse_coo_tensor(*args, **kwargs)
|
||||
elif is_compressed_sparse_layout:
|
||||
kwargs.update(layout=layout)
|
||||
yield torch.sparse_compressed_tensor(*args, **kwargs)
|
||||
else:
|
||||
assert 0 # unreachable
|
||||
return
|
||||
|
||||
def get_blockpattern(pattern, blocksize):
|
||||
basesize = pattern.shape
|
||||
assert basesize[0] % blocksize[0] == 0, (basesize, blocksize)
|
||||
assert basesize[1] % blocksize[1] == 0, (basesize, blocksize)
|
||||
blockpattern = pattern.reshape(-1,
|
||||
blocksize[0],
|
||||
basesize[1] // blocksize[1],
|
||||
blocksize[1]).transpose(-3, -2).any(-1).any(-1)
|
||||
block_ids = torch.arange(1, blockpattern.numel() + 1).reshape(blockpattern.shape)
|
||||
return (blockpattern != 0) * block_ids
|
||||
|
||||
def get_sparse_data(pattern):
|
||||
basesize = pattern.shape
|
||||
assert len(basesize) == 2, basesize # pattern is expected to be a matrix
|
||||
|
||||
# We cannot use `torch.sparse_xyz_tensor(pattern)` to
|
||||
# compute the sparse layout indices and values because
|
||||
# generate_simple_inputs is used to generate the inputs to
|
||||
# test `torch.sparse_xyz_tensor` factory functions, so
|
||||
# we'll compute the indices and values independently of
|
||||
# the factory functions.
|
||||
|
||||
indices = torch.where(pattern != 0)
|
||||
coo_indices = torch.stack(indices)
|
||||
crow_indices = torch.zeros(basesize[0] + 1, dtype=torch.int64)
|
||||
crow_indices[1:] = torch.cumsum(coo_indices[0].bincount(minlength=basesize[0]), 0)
|
||||
col_indices = coo_indices[1]
|
||||
strided_values = torch.zeros(basesize, dtype=torch.int64)
|
||||
|
||||
# the property of `values == range(1, 1+nnz)` is used in
|
||||
# get_sparse_data_with_block to relate BSR and BSC values,
|
||||
# so, don't change the following line:
|
||||
values = torch.arange(1, 1 + len(indices[0]), dtype=torch.int64)
|
||||
strided_values[indices] = values
|
||||
|
||||
indices_T = torch.where(pattern.transpose(0, 1) != 0)
|
||||
coo_indices_T = torch.stack(indices_T)
|
||||
ccol_indices = torch.zeros(basesize[1] + 1, dtype=torch.int64)
|
||||
ccol_indices[1:] = torch.cumsum(coo_indices_T[0].bincount(minlength=basesize[1]), 0)
|
||||
row_indices = coo_indices_T[1]
|
||||
csc_values = strided_values.transpose(0, 1)[indices_T]
|
||||
|
||||
return {torch.sparse_coo: (coo_indices, values),
|
||||
torch.sparse_csr: (crow_indices, col_indices, values),
|
||||
torch.sparse_csc: (ccol_indices, row_indices, csc_values),
|
||||
torch.strided: (strided_values,)}
|
||||
|
||||
def get_sparse_data_with_block(pattern, blocksize):
|
||||
nonblock_data = get_sparse_data(pattern)
|
||||
blockpattern = get_blockpattern(pattern, blocksize)
|
||||
block_data = get_sparse_data(blockpattern)
|
||||
|
||||
strided_values = nonblock_data[torch.strided][0]
|
||||
block_indices = block_data[torch.sparse_coo][0]
|
||||
bsr_values = torch.stack([strided_values[bi * blocksize[0]:(bi + 1) * blocksize[0],
|
||||
bj * blocksize[1]:(bj + 1) * blocksize[1]]
|
||||
for bi, bj in block_indices.transpose(0, 1)])
|
||||
|
||||
# here we use the property `values == range(1, 1+nnz)` and
|
||||
# `values` relation to `csc_values` (see get_sparse_data)
|
||||
# to get BSC blocks via reordering the BSR blocks:
|
||||
bsc_values = bsr_values[block_data[torch.sparse_csc][2] - 1]
|
||||
|
||||
return {torch.sparse_bsr: (*block_data[torch.sparse_csr][:2], bsr_values),
|
||||
torch.sparse_bsc: (*block_data[torch.sparse_csc][:2], bsc_values),
|
||||
**nonblock_data}
|
||||
|
||||
def get_batch_sparse_data(pattern, blocksize):
|
||||
size = pattern.shape
|
||||
if len(size) <= 2: # non-batch
|
||||
return get_sparse_data_with_block(pattern, blocksize)
|
||||
|
||||
# batch data is created recursively:
|
||||
batch_data = {}
|
||||
for i, item in enumerate(pattern):
|
||||
for layout, d in get_batch_sparse_data(item, blocksize).items():
|
||||
target = batch_data.get(layout)
|
||||
if layout is torch.sparse_coo:
|
||||
# a "batch COO" means a COO with the leading
|
||||
# sparse dimensions interpreted as batch
|
||||
# dimensions
|
||||
ext_coo_indices1 = torch.cat((torch.full((1, len(d[1])), i, dtype=torch.int64), d[0]))
|
||||
if target is None:
|
||||
target = batch_data[layout] = (ext_coo_indices1, d[1])
|
||||
else:
|
||||
target[0].set_(torch.cat((target[0], ext_coo_indices1), 1))
|
||||
target[1].set_(torch.cat((target[1], d[1])))
|
||||
else:
|
||||
if target is None:
|
||||
target = batch_data[layout] = tuple(d[j].unsqueeze(0) for j in range(len(d)))
|
||||
else:
|
||||
for j in range(len(d)):
|
||||
target[j].set_(torch.cat((target[j], d[j].unsqueeze(0))))
|
||||
return batch_data
|
||||
|
||||
def generate_values(base, densesize):
|
||||
"""Generates a tensor of shape densesize with values equal to
|
||||
|
||||
base + i_1 * 10^0 + ... + i_d * 10^{d - 1}
|
||||
|
||||
at indices i_1, ..., i_d (with 0 <= i_j < densesize[j] for any 1 <= j <=
|
||||
len(densesize))
|
||||
|
||||
This mapping produces unique values as long as
|
||||
densesize[i] < 10 for all i in range(len(densesize)).
|
||||
"""
|
||||
|
||||
if not densesize:
|
||||
return base
|
||||
if not isinstance(base, int) and base.ndim > 0:
|
||||
return torch.stack([generate_values(b, densesize) for b in base])
|
||||
if base == 0:
|
||||
return torch.zeros(densesize, dtype=torch.int64)
|
||||
r = torch.arange(densesize[0], dtype=torch.int64)
|
||||
for i, d in enumerate(densesize[1:]):
|
||||
y = torch.arange(d, dtype=torch.int64) * (10 ** (i + 1))
|
||||
r = r[..., None] + y[None, ...]
|
||||
r.add_(base)
|
||||
return r
|
||||
|
||||
if patterns is None:
|
||||
# A pattern is a 3-tuple with the following items:
|
||||
#
|
||||
# - a list of integers with the depth of two or more. The
|
||||
# integers define the sparsity patterns of the generated
|
||||
# inputs: zero values correspond to unspecified
|
||||
# elements/blocks, and non-zero values to the specified
|
||||
# elements.
|
||||
#
|
||||
# For debugging convenience, the elements with the same
|
||||
# value typically belong to the same block. However, it
|
||||
# is not a hard requirement: as long as the shape of a
|
||||
# pattern divides with block sizes, the pattern will be
|
||||
# a valid one.
|
||||
#
|
||||
# If the depth of the list is larger than two, inputs
|
||||
# with batch dimensions will be generated.
|
||||
#
|
||||
# - a list of 2-tuples of block sizes, used to generate
|
||||
# BSR/BSC tensors with various block size parameters
|
||||
#
|
||||
# - a list of tuples of dense dimensions, used to generate
|
||||
# hybrid tensors with various dense dimensions
|
||||
#
|
||||
patterns = [
|
||||
# a simple 3 x 2 tensor: non-hybrid, hybrid with 1 and 2 dense dimensions
|
||||
([[1, 2, 0],
|
||||
[1, 0, 3]], [(2, 1), (1, 3)], [(), (2,), (4, 5)]),
|
||||
# 2 x 3 batch of 3 x 2 tensors: non-hybrid and hybrid with 2 dense dimensions
|
||||
([[[[1, 2, 0],
|
||||
[1, 0, 3]],
|
||||
[[1, 2, 3],
|
||||
[1, 0, 0]],
|
||||
[[1, 0, 0],
|
||||
[1, 2, 3]]],
|
||||
[[[0, 2, 0],
|
||||
[1, 2, 3]],
|
||||
[[1, 0, 3],
|
||||
[1, 2, 0]],
|
||||
[[1, 2, 3],
|
||||
[0, 2, 0]]]], [(2, 1), (2, 3)], [(), (2,)]),
|
||||
# tensor with non-trivial blocksize
|
||||
([[0, 1, 0, 2, 0, 2],
|
||||
[0, 1, 0, 0, 2, 0],
|
||||
[3, 3, 3, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[0, 5, 0, 6, 6, 6],
|
||||
[5, 0, 5, 6, 6, 6],
|
||||
[0, 0, 0, 0, 8, 8],
|
||||
[7, 7, 7, 0, 8, 8]], [(2, 3)], [(), (4, 5)]),
|
||||
# batch tensor with variable NSE
|
||||
# Requires https://github.com/pytorch/pytorch/pull/84843 or similar.
|
||||
([[[1, 2],
|
||||
[3, 4]],
|
||||
[[1, 0],
|
||||
[0, 0]]], [(1, 1)], ([()] if enable_batch_variable_nse else []))]
|
||||
|
||||
# the main loop of the method:
|
||||
for pattern, blocksizes, densesizes in patterns:
|
||||
if not enable_hybrid:
|
||||
densesizes = [s for s in densesizes if not s]
|
||||
if not (densesizes and blocksizes):
|
||||
continue
|
||||
pattern = torch.tensor(pattern, dtype=torch.int64)
|
||||
if not enable_batch and pattern.ndim > 2:
|
||||
continue
|
||||
for blocksize in blocksizes:
|
||||
data = get_batch_sparse_data(pattern, blocksize)[layout]
|
||||
for densesize in densesizes:
|
||||
indices = [a.to(device=device, dtype=index_dtype) for a in data[:-1]]
|
||||
values = generate_values(data[-1], densesize).to(device=device, dtype=dtype)
|
||||
yield (*indices, values), dict(device=device, dtype=dtype,
|
||||
size=pattern.shape + densesize)
|
||||
|
||||
# zero-sized tensor inputs, non-batch, non-hybrid/hybrid
|
||||
if enable_zero_sized:
|
||||
for basesize, blocksizes, densesizes in [
|
||||
((2, 0), [(1, 2)], [(), (2,), (2, 3)] if enable_hybrid else [()]),
|
||||
((0, 2), [(1, 2), (2, 1), (3, 2)], [()]),
|
||||
((0, 0), [(1, 2)], [()]),
|
||||
]:
|
||||
for blocksize in blocksizes:
|
||||
for densesize in densesizes:
|
||||
if layout == torch.strided:
|
||||
indices = ()
|
||||
values = torch.empty((basesize + densesize), device=device, dtype=dtype)
|
||||
elif layout == torch.sparse_coo:
|
||||
indices = (torch.empty(len(basesize), 0, device=device, dtype=index_dtype),)
|
||||
values = torch.empty((0, *densesize), device=device, dtype=dtype)
|
||||
elif layout == torch.sparse_csr:
|
||||
crow_indices = torch.tensor([0] * (basesize[0] + 1), device=device, dtype=index_dtype)
|
||||
col_indices = torch.empty(0, device=device, dtype=index_dtype)
|
||||
indices = (crow_indices, col_indices)
|
||||
values = torch.empty((0, *densesize), device=device, dtype=dtype)
|
||||
elif layout == torch.sparse_csc:
|
||||
ccol_indices = torch.tensor([0] * (basesize[1] + 1), device=device, dtype=index_dtype)
|
||||
row_indices = torch.empty(0, device=device, dtype=index_dtype)
|
||||
indices = (ccol_indices, row_indices)
|
||||
values = torch.empty((0, *densesize), device=device, dtype=dtype)
|
||||
elif layout == torch.sparse_bsr:
|
||||
crow_indices = torch.tensor([0] * (basesize[0] // blocksize[0] + 1), device=device, dtype=index_dtype)
|
||||
col_indices = torch.empty(0, device=device, dtype=index_dtype)
|
||||
indices = (crow_indices, col_indices)
|
||||
values = torch.empty((0, *blocksize, *densesize), device=device, dtype=dtype)
|
||||
elif layout == torch.sparse_bsc:
|
||||
ccol_indices = torch.tensor([0] * (basesize[1] // blocksize[1] + 1), device=device, dtype=index_dtype)
|
||||
row_indices = torch.empty(0, device=device, dtype=index_dtype)
|
||||
indices = (ccol_indices, row_indices)
|
||||
values = torch.empty((0, *blocksize, *densesize), device=device, dtype=dtype)
|
||||
else:
|
||||
assert 0 # unreachable
|
||||
yield (*indices, values), dict(device=device, dtype=dtype, size=basesize + densesize)
|
||||
|
||||
def safeToDense(self, t):
|
||||
# coalesce is only implemented for COO
|
||||
if t.layout == torch.sparse_coo:
|
||||
|
|
|
|||
Loading…
Reference in a new issue