def add_nnz_(src: SparseTensor, other: torch.Tensor, layout: Optional[str] = None) -> SparseTensor: value = src.storage.value() if value is not None: value = value.add_(other.to(value.dtype)) else: value = other.add(1) return src.set_value_(value, layout=layout)
def add_(src: SparseTensor, other: torch.Tensor) -> SparseTensor: rowptr, col, value = src.csr() if other.size(0) == src.size(0) and other.size(1) == 1: # Row-wise. other = gather_csr(other.squeeze(1), rowptr) elif other.size(0) == 1 and other.size(1) == src.size(1): # Col-wise. other = other.squeeze(0)[col] else: raise ValueError( f'Size mismatch: Expected size ({src.size(0)}, 1, ...) or ' f'(1, {src.size(1)}, ...), but got size {other.size()}.') if value is not None: value = value.add_(other.to(value.dtype)) else: value = other.add_(1) return src.set_value_(value, layout='coo')
def test_cat(device): row, col = tensor([[0, 0, 1], [0, 1, 2]], torch.long, device) mat1 = SparseTensor(row=row, col=col) mat1.fill_cache_() row, col = tensor([[0, 0, 1, 2], [0, 1, 1, 0]], torch.long, device) mat2 = SparseTensor(row=row, col=col) mat2.fill_cache_() out = cat([mat1, mat2], dim=0) assert out.to_dense().tolist() == [[1, 1, 0], [0, 0, 1], [1, 1, 0], [0, 1, 0], [1, 0, 0]] assert out.storage.has_row() assert out.storage.has_rowptr() assert out.storage.has_rowcount() assert out.storage.num_cached_keys() == 1 out = cat([mat1, mat2], dim=1) assert out.to_dense().tolist() == [[1, 1, 0, 1, 1], [0, 0, 1, 0, 1], [0, 0, 0, 1, 0]] assert out.storage.has_row() assert not out.storage.has_rowptr() assert out.storage.num_cached_keys() == 2 out = cat([mat1, mat2], dim=(0, 1)) assert out.to_dense().tolist() == [[1, 1, 0, 0, 0], [0, 0, 1, 0, 0], [0, 0, 0, 1, 1], [0, 0, 0, 0, 1], [0, 0, 0, 1, 0]] assert out.storage.has_row() assert out.storage.has_rowptr() assert out.storage.num_cached_keys() == 5 value = torch.randn((mat1.nnz(), 4), device=device) mat1 = mat1.set_value_(value, layout='coo') out = cat([mat1, mat1], dim=-1) assert out.storage.value().size() == (mat1.nnz(), 8) assert out.storage.has_row() assert out.storage.has_rowptr() assert out.storage.num_cached_keys() == 5