storage.py 23.1 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
import warnings
rusty1s's avatar
rusty1s committed
2
from typing import Optional, List, Tuple
rusty1s's avatar
rusty1s committed
3

rusty1s's avatar
rusty1s committed
4
import torch
rusty1s's avatar
rusty1s committed
5
from torch_scatter import segment_csr, scatter_add
rusty1s's avatar
rusty1s committed
6
from torch_sparse.utils import Final
rusty1s's avatar
rusty1s committed
7

rusty1s's avatar
rusty1s committed
8
layouts: Final[List[str]] = ['coo', 'csr', 'csc']
rusty1s's avatar
rusty1s committed
9
10


rusty1s's avatar
rusty1s committed
11
def get_layout(layout: Optional[str] = None) -> str:
rusty1s's avatar
rusty1s committed
12
13
14
15
    if layout is None:
        layout = 'coo'
        warnings.warn('`layout` argument unset, using default layout '
                      '"coo". This may lead to unexpected behaviour.')
rusty1s's avatar
rusty1s committed
16
    assert layout == 'coo' or layout == 'csr' or layout == 'csc'
rusty1s's avatar
rusty1s committed
17
18
19
    return layout


rusty1s's avatar
rusty1s committed
20
@torch.jit.script
rusty1s's avatar
rusty1s committed
21
class SparseStorage(object):
rusty1s's avatar
rusty1s committed
22
23
24
25
    _row: Optional[torch.Tensor]
    _rowptr: Optional[torch.Tensor]
    _col: torch.Tensor
    _value: Optional[torch.Tensor]
rusty1s's avatar
rusty1s committed
26
    _sparse_sizes: Tuple[int, int]
rusty1s's avatar
rusty1s committed
27
28
29
30
31
32
    _rowcount: Optional[torch.Tensor]
    _colptr: Optional[torch.Tensor]
    _colcount: Optional[torch.Tensor]
    _csr2csc: Optional[torch.Tensor]
    _csc2csr: Optional[torch.Tensor]

33
    def __init__(self, row: Optional[torch.Tensor] = None,
rusty1s's avatar
rusty1s committed
34
35
36
                 rowptr: Optional[torch.Tensor] = None,
                 col: Optional[torch.Tensor] = None,
                 value: Optional[torch.Tensor] = None,
rusty1s's avatar
rusty1s committed
37
                 sparse_sizes: Optional[Tuple[int, int]] = None,
rusty1s's avatar
rusty1s committed
38
39
40
41
42
43
                 rowcount: Optional[torch.Tensor] = None,
                 colptr: Optional[torch.Tensor] = None,
                 colcount: Optional[torch.Tensor] = None,
                 csr2csc: Optional[torch.Tensor] = None,
                 csc2csr: Optional[torch.Tensor] = None,
                 is_sorted: bool = False):
rusty1s's avatar
rusty1s committed
44

rusty1s's avatar
rusty1s committed
45
46
47
48
        assert row is not None or rowptr is not None
        assert col is not None
        assert col.dtype == torch.long
        assert col.dim() == 1
rusty1s's avatar
rusty1s committed
49
        col = col.contiguous()
rusty1s's avatar
rusty1s committed
50

rusty1s's avatar
rusty1s committed
51
        if sparse_sizes is None:
rusty1s's avatar
rusty1s committed
52
53
            if rowptr is not None:
                M = rowptr.numel() - 1
rusty1s's avatar
rusty1s committed
54
            elif row is not None and row.numel() > 0:
rusty1s's avatar
rusty1s committed
55
                M = row.max().item() + 1
rusty1s's avatar
rusty1s committed
56
57
            elif row is not None and row.numel() == 0:
                M = 0
rusty1s's avatar
rusty1s committed
58
59
            else:
                raise ValueError
rusty1s's avatar
rusty1s committed
60
61
62
63
            if col.numel() > 0:
                N = col.max().item() + 1
            else:
                N = 0
rusty1s's avatar
rusty1s committed
64
            sparse_sizes = (int(M), int(N))
rusty1s's avatar
rusty1s committed
65
        else:
rusty1s's avatar
rusty1s committed
66
            assert len(sparse_sizes) == 2
rusty1s's avatar
rusty1s committed
67
            if row is not None and row.numel() > 0:
rusty1s's avatar
update  
rusty1s committed
68
                assert row.max().item() < sparse_sizes[0]
rusty1s's avatar
rusty1s committed
69
70
            if col.numel() > 0:
                assert col.max().item() < sparse_sizes[1]
rusty1s's avatar
rusty1s committed
71

rusty1s's avatar
rusty1s committed
72
73
74
75
76
        if row is not None:
            assert row.dtype == torch.long
            assert row.device == col.device
            assert row.dim() == 1
            assert row.numel() == col.numel()
rusty1s's avatar
rusty1s committed
77
            row = row.contiguous()
rusty1s's avatar
rusty1s committed
78

rusty1s's avatar
rusty1s committed
79
        if rowptr is not None:
rusty1s's avatar
rusty1s committed
80
            assert rowptr.dtype == torch.long
rusty1s's avatar
rusty1s committed
81
82
            assert rowptr.device == col.device
            assert rowptr.dim() == 1
rusty1s's avatar
rusty1s committed
83
            assert rowptr.numel() - 1 == sparse_sizes[0]
rusty1s's avatar
rusty1s committed
84
            rowptr = rowptr.contiguous()
rusty1s's avatar
rusty1s committed
85

rusty1s's avatar
rusty1s committed
86
87
88
        if value is not None:
            assert value.device == col.device
            assert value.size(0) == col.size(0)
rusty1s's avatar
rusty1s committed
89
            value = value.contiguous()
rusty1s's avatar
rusty1s committed
90
91
92
93
94

        if rowcount is not None:
            assert rowcount.dtype == torch.long
            assert rowcount.device == col.device
            assert rowcount.dim() == 1
rusty1s's avatar
rusty1s committed
95
            assert rowcount.numel() == sparse_sizes[0]
rusty1s's avatar
rusty1s committed
96
            rowcount = rowcount.contiguous()
rusty1s's avatar
rusty1s committed
97

rusty1s's avatar
rusty1s committed
98
        if colptr is not None:
rusty1s's avatar
rusty1s committed
99
            assert colptr.dtype == torch.long
rusty1s's avatar
rusty1s committed
100
101
            assert colptr.device == col.device
            assert colptr.dim() == 1
rusty1s's avatar
rusty1s committed
102
            assert colptr.numel() - 1 == sparse_sizes[1]
rusty1s's avatar
rusty1s committed
103
            colptr = colptr.contiguous()
rusty1s's avatar
rusty1s committed
104
105
106
107
108

        if colcount is not None:
            assert colcount.dtype == torch.long
            assert colcount.device == col.device
            assert colcount.dim() == 1
rusty1s's avatar
rusty1s committed
109
            assert colcount.numel() == sparse_sizes[1]
rusty1s's avatar
rusty1s committed
110
            colcount = colcount.contiguous()
rusty1s's avatar
rusty1s committed
111

rusty1s's avatar
rusty1s committed
112
113
        if csr2csc is not None:
            assert csr2csc.dtype == torch.long
rusty1s's avatar
rusty1s committed
114
            assert csr2csc.device == col.device
rusty1s's avatar
rusty1s committed
115
            assert csr2csc.dim() == 1
rusty1s's avatar
rusty1s committed
116
            assert csr2csc.numel() == col.size(0)
rusty1s's avatar
rusty1s committed
117
            csr2csc = csr2csc.contiguous()
rusty1s's avatar
rusty1s committed
118

rusty1s's avatar
rusty1s committed
119
120
        if csc2csr is not None:
            assert csc2csr.dtype == torch.long
rusty1s's avatar
rusty1s committed
121
            assert csc2csr.device == col.device
rusty1s's avatar
rusty1s committed
122
            assert csc2csr.dim() == 1
rusty1s's avatar
rusty1s committed
123
            assert csc2csr.numel() == col.size(0)
rusty1s's avatar
rusty1s committed
124
            csc2csr = csc2csr.contiguous()
rusty1s's avatar
rusty1s committed
125

rusty1s's avatar
rusty1s committed
126
127
128
        self._row = row
        self._rowptr = rowptr
        self._col = col
rusty1s's avatar
rusty1s committed
129
        self._value = value
rusty1s's avatar
rusty1s committed
130
        self._sparse_sizes = tuple(sparse_sizes)
rusty1s's avatar
rusty1s committed
131
        self._rowcount = rowcount
rusty1s's avatar
rusty1s committed
132
        self._colptr = colptr
rusty1s's avatar
rusty1s committed
133
        self._colcount = colcount
rusty1s's avatar
rusty1s committed
134
135
        self._csr2csc = csr2csc
        self._csc2csr = csc2csr
rusty1s's avatar
rusty1s committed
136

rusty1s's avatar
rusty1s committed
137
        if not is_sorted:
rusty1s's avatar
rusty1s committed
138
            idx = self._col.new_zeros(self._col.numel() + 1)
rusty1s's avatar
rusty1s committed
139
140
141
            idx[1:] = self.row()
            idx[1:] *= self._sparse_sizes[1]
            idx[1:] += self._col
rusty1s's avatar
rusty1s committed
142
            if (idx[1:] < idx[:-1]).any():
rusty1s's avatar
rusty1s committed
143
                perm = idx[1:].argsort()
rusty1s's avatar
rusty1s committed
144
                self._row = self.row()[perm]
rusty1s's avatar
rusty1s committed
145
                self._col = self._col[perm]
rusty1s's avatar
rusty1s committed
146
147
                if value is not None:
                    self._value = value[perm]
rusty1s's avatar
rusty1s committed
148
149
150
                self._csr2csc = None
                self._csc2csr = None

rusty1s's avatar
rusty1s committed
151
152
    @classmethod
    def empty(self):
rusty1s's avatar
rusty1s committed
153
154
155
156
157
158
        row = torch.tensor([], dtype=torch.long)
        col = torch.tensor([], dtype=torch.long)
        return SparseStorage(row=row, rowptr=None, col=col, value=None,
                             sparse_sizes=(0, 0), rowcount=None, colptr=None,
                             colcount=None, csr2csc=None, csc2csr=None,
                             is_sorted=True)
rusty1s's avatar
rusty1s committed
159

rusty1s's avatar
rusty1s committed
160
    def has_row(self) -> bool:
rusty1s's avatar
rusty1s committed
161
        return self._row is not None
rusty1s's avatar
rusty1s committed
162

rusty1s's avatar
rusty1s committed
163
    def row(self):
rusty1s's avatar
rusty1s committed
164
165
166
        row = self._row
        if row is not None:
            return row
rusty1s's avatar
rusty1s committed
167

rusty1s's avatar
rusty1s committed
168
169
        rowptr = self._rowptr
        if rowptr is not None:
rusty1s's avatar
rusty1s committed
170
            row = torch.ops.torch_sparse.ptr2ind(rowptr, self._col.numel())
rusty1s's avatar
rusty1s committed
171
172
173
174
175
176
            self._row = row
            return row

        raise ValueError

    def has_rowptr(self) -> bool:
rusty1s's avatar
rusty1s committed
177
178
        return self._rowptr is not None

rusty1s's avatar
rusty1s committed
179
180
181
182
    def rowptr(self) -> torch.Tensor:
        rowptr = self._rowptr
        if rowptr is not None:
            return rowptr
rusty1s's avatar
rusty1s committed
183

rusty1s's avatar
rusty1s committed
184
185
        row = self._row
        if row is not None:
rusty1s's avatar
rusty1s committed
186
            rowptr = torch.ops.torch_sparse.ind2ptr(row, self._sparse_sizes[0])
rusty1s's avatar
rusty1s committed
187
188
189
190
191
192
            self._rowptr = rowptr
            return rowptr

        raise ValueError

    def col(self) -> torch.Tensor:
rusty1s's avatar
rusty1s committed
193
        return self._col
rusty1s's avatar
rusty1s committed
194

rusty1s's avatar
rusty1s committed
195
    def has_value(self) -> bool:
rusty1s's avatar
rusty1s committed
196
        return self._value is not None
rusty1s's avatar
rusty1s committed
197

rusty1s's avatar
rusty1s committed
198
    def value(self) -> Optional[torch.Tensor]:
rusty1s's avatar
rusty1s committed
199
200
        return self._value

201
    def set_value_(self, value: Optional[torch.Tensor],
rusty1s's avatar
rusty1s committed
202
203
                   layout: Optional[str] = None):
        if value is not None:
rusty1s's avatar
rusty1s committed
204
            if get_layout(layout) == 'csc':
rusty1s's avatar
rusty1s committed
205
206
207
208
                value = value[self.csc2csr()]
            value = value.contiguous()
            assert value.device == self._col.device
            assert value.size(0) == self._col.numel()
rusty1s's avatar
rusty1s committed
209

rusty1s's avatar
rusty1s committed
210
211
        self._value = value
        return self
rusty1s's avatar
rusty1s committed
212

213
    def set_value(self, value: Optional[torch.Tensor],
rusty1s's avatar
rusty1s committed
214
215
                  layout: Optional[str] = None):
        if value is not None:
rusty1s's avatar
rusty1s committed
216
            if get_layout(layout) == 'csc':
rusty1s's avatar
rusty1s committed
217
218
219
220
                value = value[self.csc2csr()]
            value = value.contiguous()
            assert value.device == self._col.device
            assert value.size(0) == self._col.numel()
rusty1s's avatar
rusty1s committed
221

222
223
224
225
226
        return SparseStorage(row=self._row, rowptr=self._rowptr, col=self._col,
                             value=value, sparse_sizes=self._sparse_sizes,
                             rowcount=self._rowcount, colptr=self._colptr,
                             colcount=self._colcount, csr2csc=self._csr2csc,
                             csc2csr=self._csc2csr, is_sorted=True)
rusty1s's avatar
rusty1s committed
227

rusty1s's avatar
rusty1s committed
228
    def sparse_sizes(self) -> Tuple[int, int]:
rusty1s's avatar
rusty1s committed
229
        return self._sparse_sizes
rusty1s's avatar
rusty1s committed
230

rusty1s's avatar
rusty1s committed
231
232
    def sparse_size(self, dim: int) -> int:
        return self._sparse_sizes[dim]
rusty1s's avatar
rusty1s committed
233

rusty1s's avatar
rusty1s committed
234
    def sparse_resize(self, sparse_sizes: Tuple[int, int]):
rusty1s's avatar
rusty1s committed
235
236
        assert len(sparse_sizes) == 2
        old_sparse_sizes, nnz = self._sparse_sizes, self._col.numel()
rusty1s's avatar
rusty1s committed
237

rusty1s's avatar
rusty1s committed
238
        diff_0 = sparse_sizes[0] - old_sparse_sizes[0]
rusty1s's avatar
rusty1s committed
239
240
241
242
243
244
        rowcount, rowptr = self._rowcount, self._rowptr
        if diff_0 > 0:
            if rowptr is not None:
                rowptr = torch.cat([rowptr, rowptr.new_full((diff_0, ), nnz)])
            if rowcount is not None:
                rowcount = torch.cat([rowcount, rowcount.new_zeros(diff_0)])
rusty1s's avatar
rusty1s committed
245
        elif diff_0 < 0:
rusty1s's avatar
rusty1s committed
246
247
248
249
250
            if rowptr is not None:
                rowptr = rowptr[:-diff_0]
            if rowcount is not None:
                rowcount = rowcount[:-diff_0]

rusty1s's avatar
rusty1s committed
251
        diff_1 = sparse_sizes[1] - old_sparse_sizes[1]
rusty1s's avatar
rusty1s committed
252
253
254
255
256
257
        colcount, colptr = self._colcount, self._colptr
        if diff_1 > 0:
            if colptr is not None:
                colptr = torch.cat([colptr, colptr.new_full((diff_1, ), nnz)])
            if colcount is not None:
                colcount = torch.cat([colcount, colcount.new_zeros(diff_1)])
rusty1s's avatar
rusty1s committed
258
        elif diff_1 < 0:
rusty1s's avatar
rusty1s committed
259
260
261
262
263
            if colptr is not None:
                colptr = colptr[:-diff_1]
            if colcount is not None:
                colcount = colcount[:-diff_1]

264
265
266
267
268
        return SparseStorage(row=self._row, rowptr=rowptr, col=self._col,
                             value=self._value, sparse_sizes=sparse_sizes,
                             rowcount=rowcount, colptr=colptr,
                             colcount=colcount, csr2csc=self._csr2csc,
                             csc2csr=self._csc2csr, is_sorted=True)
rusty1s's avatar
rusty1s committed
269

rusty1s's avatar
rusty1s committed
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
    def sparse_reshape(self, num_rows: int, num_cols: int):
        assert num_rows > 0 or num_rows == -1
        assert num_cols > 0 or num_cols == -1
        assert num_rows > 0 or num_cols > 0

        total = self.sparse_size(0) * self.sparse_size(1)

        if num_rows == -1:
            num_rows = total // num_cols

        if num_cols == -1:
            num_cols = total // num_rows

        assert num_rows * num_cols == total

        idx = self.sparse_size(1) * self.row() + self.col()

rusty1s's avatar
rusty1s committed
287
        row = idx // num_cols
rusty1s's avatar
rusty1s committed
288
        col = idx % num_cols
rusty1s's avatar
rusty1s committed
289
        assert row.dtype == torch.long and col.dtype == torch.long
rusty1s's avatar
rusty1s committed
290
291
292
293
294
295

        return SparseStorage(row=row, rowptr=None, col=col, value=self._value,
                             sparse_sizes=(num_rows, num_cols), rowcount=None,
                             colptr=None, colcount=None, csr2csc=None,
                             csc2csr=None, is_sorted=True)

rusty1s's avatar
rusty1s committed
296
    def has_rowcount(self) -> bool:
rusty1s's avatar
rusty1s committed
297
298
        return self._rowcount is not None

rusty1s's avatar
rusty1s committed
299
300
301
302
303
304
    def rowcount(self) -> torch.Tensor:
        rowcount = self._rowcount
        if rowcount is not None:
            return rowcount

        rowptr = self.rowptr()
305
        rowcount = rowptr[1:] - rowptr[:-1]
rusty1s's avatar
rusty1s committed
306
307
        self._rowcount = rowcount
        return rowcount
rusty1s's avatar
rusty1s committed
308

rusty1s's avatar
rusty1s committed
309
    def has_colptr(self) -> bool:
rusty1s's avatar
rusty1s committed
310
        return self._colptr is not None
rusty1s's avatar
rusty1s committed
311

rusty1s's avatar
rusty1s committed
312
313
314
    def colptr(self) -> torch.Tensor:
        colptr = self._colptr
        if colptr is not None:
rusty1s's avatar
rusty1s committed
315
            return colptr
rusty1s's avatar
rusty1s committed
316

rusty1s's avatar
rusty1s committed
317
318
        csr2csc = self._csr2csc
        if csr2csc is not None:
rusty1s's avatar
rusty1s committed
319
320
            colptr = torch.ops.torch_sparse.ind2ptr(self._col[csr2csc],
                                                    self._sparse_sizes[1])
rusty1s's avatar
rusty1s committed
321
        else:
rusty1s's avatar
rusty1s committed
322
            colptr = self._col.new_zeros(self._sparse_sizes[1] + 1)
rusty1s's avatar
rusty1s committed
323
324
325
326
327
            torch.cumsum(self.colcount(), dim=0, out=colptr[1:])
        self._colptr = colptr
        return colptr

    def has_colcount(self) -> bool:
rusty1s's avatar
rusty1s committed
328
329
        return self._colcount is not None

rusty1s's avatar
rusty1s committed
330
331
332
333
334
335
336
    def colcount(self) -> torch.Tensor:
        colcount = self._colcount
        if colcount is not None:
            return colcount

        colptr = self._colptr
        if colptr is not None:
337
            colcount = colptr[1:] - colptr[:-1]
rusty1s's avatar
rusty1s committed
338
        else:
339
340
            colcount = scatter_add(torch.ones_like(self._col), self._col,
                                   dim_size=self._sparse_sizes[1])
rusty1s's avatar
rusty1s committed
341
342
        self._colcount = colcount
        return colcount
rusty1s's avatar
rusty1s committed
343

rusty1s's avatar
rusty1s committed
344
    def has_csr2csc(self) -> bool:
rusty1s's avatar
rusty1s committed
345
        return self._csr2csc is not None
rusty1s's avatar
rusty1s committed
346

rusty1s's avatar
rusty1s committed
347
348
349
350
    def csr2csc(self) -> torch.Tensor:
        csr2csc = self._csr2csc
        if csr2csc is not None:
            return csr2csc
rusty1s's avatar
rusty1s committed
351

rusty1s's avatar
rusty1s committed
352
        idx = self._sparse_sizes[0] * self._col + self.row()
rusty1s's avatar
rusty1s committed
353
354
355
356
357
        csr2csc = idx.argsort()
        self._csr2csc = csr2csc
        return csr2csc

    def has_csc2csr(self) -> bool:
rusty1s's avatar
rusty1s committed
358
359
        return self._csc2csr is not None

rusty1s's avatar
rusty1s committed
360
361
362
363
    def csc2csr(self) -> torch.Tensor:
        csc2csr = self._csc2csr
        if csc2csr is not None:
            return csc2csr
rusty1s's avatar
rusty1s committed
364

rusty1s's avatar
rusty1s committed
365
366
367
        csc2csr = self.csr2csc().argsort()
        self._csc2csr = csc2csr
        return csc2csr
rusty1s's avatar
rusty1s committed
368

rusty1s's avatar
rusty1s committed
369
370
    def is_coalesced(self) -> bool:
        idx = self._col.new_full((self._col.numel() + 1, ), -1)
rusty1s's avatar
rusty1s committed
371
        idx[1:] = self._sparse_sizes[1] * self.row() + self._col
rusty1s's avatar
rusty1s committed
372
373
374
375
        return bool((idx[1:] > idx[:-1]).all())

    def coalesce(self, reduce: str = "add"):
        idx = self._col.new_full((self._col.numel() + 1, ), -1)
rusty1s's avatar
rusty1s committed
376
        idx[1:] = self._sparse_sizes[1] * self.row() + self._col
rusty1s's avatar
rusty1s committed
377
        mask = idx[1:] > idx[:-1]
rusty1s's avatar
rusty1s committed
378

rusty1s's avatar
rusty1s committed
379
        if mask.all():  # Skip if indices are already coalesced.
rusty1s's avatar
rusty1s committed
380
381
            return self

rusty1s's avatar
rusty1s committed
382
383
        row = self.row()[mask]
        col = self._col[mask]
rusty1s's avatar
rusty1s committed
384

rusty1s's avatar
rusty1s committed
385
386
        value = self._value
        if value is not None:
rusty1s's avatar
reset  
rusty1s committed
387
            ptr = mask.nonzero().flatten()
rusty1s's avatar
rusty1s committed
388
            ptr = torch.cat([ptr, ptr.new_full((1, ), value.size(0))])
rusty1s's avatar
rusty1s committed
389
            value = segment_csr(value, ptr, reduce=reduce)
rusty1s's avatar
rusty1s committed
390

391
392
393
394
        return SparseStorage(row=row, rowptr=None, col=col, value=value,
                             sparse_sizes=self._sparse_sizes, rowcount=None,
                             colptr=None, colcount=None, csr2csc=None,
                             csc2csr=None, is_sorted=True)
rusty1s's avatar
rusty1s committed
395
396
397
398
399
400
401
402
403

    def fill_cache_(self):
        self.row()
        self.rowptr()
        self.rowcount()
        self.colptr()
        self.colcount()
        self.csr2csc()
        self.csc2csr()
rusty1s's avatar
rusty1s committed
404
        return self
rusty1s's avatar
rusty1s committed
405

rusty1s's avatar
rusty1s committed
406
407
408
409
410
411
    def clear_cache_(self):
        self._rowcount = None
        self._colptr = None
        self._colcount = None
        self._csr2csc = None
        self._csc2csr = None
rusty1s's avatar
rusty1s committed
412
        return self
rusty1s's avatar
rusty1s committed
413

rusty1s's avatar
rusty1s committed
414
415
    def cached_keys(self) -> List[str]:
        keys: List[str] = []
rusty1s's avatar
rusty1s committed
416
        if self.has_rowcount():
rusty1s's avatar
rusty1s committed
417
            keys.append('rowcount')
rusty1s's avatar
rusty1s committed
418
        if self.has_colptr():
rusty1s's avatar
rusty1s committed
419
            keys.append('colptr')
rusty1s's avatar
rusty1s committed
420
        if self.has_colcount():
rusty1s's avatar
rusty1s committed
421
            keys.append('colcount')
rusty1s's avatar
rusty1s committed
422
        if self.has_csr2csc():
rusty1s's avatar
rusty1s committed
423
            keys.append('csr2csc')
rusty1s's avatar
rusty1s committed
424
        if self.has_csc2csr():
rusty1s's avatar
rusty1s committed
425
426
427
428
429
            keys.append('csc2csr')
        return keys

    def num_cached_keys(self) -> int:
        return len(self.cached_keys())
rusty1s's avatar
rusty1s committed
430

rusty1s's avatar
rusty1s committed
431
    def copy(self):
432
433
434
435
436
437
        return SparseStorage(row=self._row, rowptr=self._rowptr, col=self._col,
                             value=self._value,
                             sparse_sizes=self._sparse_sizes,
                             rowcount=self._rowcount, colptr=self._colptr,
                             colcount=self._colcount, csr2csc=self._csr2csc,
                             csc2csr=self._csc2csr, is_sorted=True)
rusty1s's avatar
rusty1s committed
438

rusty1s's avatar
test  
rusty1s committed
439
    def clone(self):
rusty1s's avatar
rusty1s committed
440
441
442
443
444
445
        row = self._row
        if row is not None:
            row = row.clone()
        rowptr = self._rowptr
        if rowptr is not None:
            rowptr = rowptr.clone()
rusty1s's avatar
rusty1s committed
446
        col = self._col.clone()
rusty1s's avatar
rusty1s committed
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
        value = self._value
        if value is not None:
            value = value.clone()
        rowcount = self._rowcount
        if rowcount is not None:
            rowcount = rowcount.clone()
        colptr = self._colptr
        if colptr is not None:
            colptr = colptr.clone()
        colcount = self._colcount
        if colcount is not None:
            colcount = colcount.clone()
        csr2csc = self._csr2csc
        if csr2csc is not None:
            csr2csc = csr2csc.clone()
        csc2csr = self._csc2csr
        if csc2csr is not None:
            csc2csr = csc2csr.clone()
rusty1s's avatar
rusty1s committed
465

466
467
468
469
470
        return SparseStorage(row=row, rowptr=rowptr, col=col, value=value,
                             sparse_sizes=self._sparse_sizes,
                             rowcount=rowcount, colptr=colptr,
                             colcount=colcount, csr2csc=csr2csc,
                             csc2csr=csc2csr, is_sorted=True)
rusty1s's avatar
rusty1s committed
471

rusty1s's avatar
rusty1s committed
472
    def type_as(self, tensor: torch.Tensor):
rusty1s's avatar
rusty1s committed
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
        value = self._value
        if value is not None:
            if tensor.dtype == value.dtype:
                return self
            else:
                return self.set_value(value.type_as(tensor), layout='coo')
        else:
            return self

    def device_as(self, tensor: torch.Tensor, non_blocking: bool = False):
        if tensor.device == self._col.device:
            return self

        row = self._row
        if row is not None:
            row = row.to(tensor.device, non_blocking=non_blocking)
        rowptr = self._rowptr
        if rowptr is not None:
            rowptr = rowptr.to(tensor.device, non_blocking=non_blocking)
        col = self._col.to(tensor.device, non_blocking=non_blocking)
        value = self._value
        if value is not None:
            value = value.to(tensor.device, non_blocking=non_blocking)
        rowcount = self._rowcount
        if rowcount is not None:
            rowcount = rowcount.to(tensor.device, non_blocking=non_blocking)
        colptr = self._colptr
        if colptr is not None:
            colptr = colptr.to(tensor.device, non_blocking=non_blocking)
        colcount = self._colcount
        if colcount is not None:
            colcount = colcount.to(tensor.device, non_blocking=non_blocking)
        csr2csc = self._csr2csc
        if csr2csc is not None:
            csr2csc = csr2csc.to(tensor.device, non_blocking=non_blocking)
        csc2csr = self._csc2csr
        if csc2csr is not None:
            csc2csr = csc2csr.to(tensor.device, non_blocking=non_blocking)
rusty1s's avatar
rusty1s committed
511

512
513
514
515
516
        return SparseStorage(row=row, rowptr=rowptr, col=col, value=value,
                             sparse_sizes=self._sparse_sizes,
                             rowcount=rowcount, colptr=colptr,
                             colcount=colcount, csr2csc=csr2csc,
                             csc2csr=csc2csr, is_sorted=True)
rusty1s's avatar
rusty1s committed
517

rusty1s's avatar
rusty1s committed
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
    def cuda(self):
        new_col = self._col.cuda()
        if new_col.device == self._col.device:
            return self

        row = self._row
        if row is not None:
            row = row.cuda()
        rowptr = self._rowptr
        if rowptr is not None:
            rowptr = rowptr.cuda()
        value = self._value
        if value is not None:
            value = value.cuda()
        rowcount = self._rowcount
        if rowcount is not None:
            rowcount = rowcount.cuda()
        colptr = self._colptr
        if colptr is not None:
            colptr = colptr.cuda()
        colcount = self._colcount
        if colcount is not None:
            colcount = colcount.cuda()
        csr2csc = self._csr2csc
        if csr2csc is not None:
            csr2csc = csr2csc.cuda()
        csc2csr = self._csc2csr
        if csc2csr is not None:
            csc2csr = csc2csr.cuda()

        return SparseStorage(row=row, rowptr=rowptr, col=new_col, value=value,
                             sparse_sizes=self._sparse_sizes,
                             rowcount=rowcount, colptr=colptr,
                             colcount=colcount, csr2csc=csr2csc,
                             csc2csr=csc2csr, is_sorted=True)

rusty1s's avatar
rusty1s committed
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
    def pin_memory(self):
        row = self._row
        if row is not None:
            row = row.pin_memory()
        rowptr = self._rowptr
        if rowptr is not None:
            rowptr = rowptr.pin_memory()
        col = self._col.pin_memory()
        value = self._value
        if value is not None:
            value = value.pin_memory()
        rowcount = self._rowcount
        if rowcount is not None:
            rowcount = rowcount.pin_memory()
        colptr = self._colptr
        if colptr is not None:
            colptr = colptr.pin_memory()
        colcount = self._colcount
        if colcount is not None:
            colcount = colcount.pin_memory()
        csr2csc = self._csr2csc
        if csr2csc is not None:
            csr2csc = csr2csc.pin_memory()
        csc2csr = self._csc2csr
        if csc2csr is not None:
            csc2csr = csc2csr.pin_memory()
rusty1s's avatar
rusty1s committed
580

581
582
583
584
585
        return SparseStorage(row=row, rowptr=rowptr, col=col, value=value,
                             sparse_sizes=self._sparse_sizes,
                             rowcount=rowcount, colptr=colptr,
                             colcount=colcount, csr2csc=csr2csc,
                             csc2csr=csc2csr, is_sorted=True)
rusty1s's avatar
rusty1s committed
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676

    def is_pinned(self) -> bool:
        is_pinned = True
        row = self._row
        if row is not None:
            is_pinned = is_pinned and row.is_pinned()
        rowptr = self._rowptr
        if rowptr is not None:
            is_pinned = is_pinned and rowptr.is_pinned()
        is_pinned = self._col.is_pinned()
        value = self._value
        if value is not None:
            is_pinned = is_pinned and value.is_pinned()
        rowcount = self._rowcount
        if rowcount is not None:
            is_pinned = is_pinned and rowcount.is_pinned()
        colptr = self._colptr
        if colptr is not None:
            is_pinned = is_pinned and colptr.is_pinned()
        colcount = self._colcount
        if colcount is not None:
            is_pinned = is_pinned and colcount.is_pinned()
        csr2csc = self._csr2csc
        if csr2csc is not None:
            is_pinned = is_pinned and csr2csc.is_pinned()
        csc2csr = self._csc2csr
        if csc2csr is not None:
            is_pinned = is_pinned and csc2csr.is_pinned()
        return is_pinned


def share_memory_(self) -> SparseStorage:
    row = self._row
    if row is not None:
        row.share_memory_()
    rowptr = self._rowptr
    if rowptr is not None:
        rowptr.share_memory_()
    self._col.share_memory_()
    value = self._value
    if value is not None:
        value.share_memory_()
    rowcount = self._rowcount
    if rowcount is not None:
        rowcount.share_memory_()
    colptr = self._colptr
    if colptr is not None:
        colptr.share_memory_()
    colcount = self._colcount
    if colcount is not None:
        colcount.share_memory_()
    csr2csc = self._csr2csc
    if csr2csc is not None:
        csr2csc.share_memory_()
    csc2csr = self._csc2csr
    if csc2csr is not None:
        csc2csr.share_memory_()


def is_shared(self) -> bool:
    is_shared = True
    row = self._row
    if row is not None:
        is_shared = is_shared and row.is_shared()
    rowptr = self._rowptr
    if rowptr is not None:
        is_shared = is_shared and rowptr.is_shared()
    is_shared = is_shared and self._col.is_shared()
    value = self._value
    if value is not None:
        is_shared = is_shared and value.is_shared()
    rowcount = self._rowcount
    if rowcount is not None:
        is_shared = is_shared and rowcount.is_shared()
    colptr = self._colptr
    if colptr is not None:
        is_shared = is_shared and colptr.is_shared()
    colcount = self._colcount
    if colcount is not None:
        is_shared = is_shared and colcount.is_shared()
    csr2csc = self._csr2csc
    if csr2csc is not None:
        is_shared = is_shared and csr2csc.is_shared()
    csc2csr = self._csc2csr
    if csc2csr is not None:
        is_shared = is_shared and csc2csr.is_shared()
    return is_shared


SparseStorage.share_memory_ = share_memory_
SparseStorage.is_shared = is_shared