storage.py 23.9 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
update  
rusty1s committed
37
38
                 sparse_sizes: Optional[Tuple[Optional[int],
                                              Optional[int]]] = None,
rusty1s's avatar
rusty1s committed
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,
44
45
                 is_sorted: bool = False,
                 trust_data: bool = False):
rusty1s's avatar
rusty1s committed
46

rusty1s's avatar
rusty1s committed
47
48
49
50
        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
51
        col = col.contiguous()
rusty1s's avatar
rusty1s committed
52

rusty1s's avatar
update  
rusty1s committed
53
54
        M: int = 0
        if sparse_sizes is None or sparse_sizes[0] is None:
rusty1s's avatar
rusty1s committed
55
56
            if rowptr is not None:
                M = rowptr.numel() - 1
rusty1s's avatar
rusty1s committed
57
            elif row is not None and row.numel() > 0:
rusty1s's avatar
update  
rusty1s committed
58
59
                M = int(row.max()) + 1
        else:
rusty1s's avatar
update  
rusty1s committed
60
61
62
            _M = sparse_sizes[0]
            assert _M is not None
            M = _M
rusty1s's avatar
update  
rusty1s committed
63
            if rowptr is not None:
rusty1s's avatar
update  
rusty1s committed
64
                assert rowptr.numel() - 1 == M
rusty1s's avatar
update  
rusty1s committed
65
            elif row is not None and row.numel() > 0:
66
                assert trust_data or int(row.max()) < M
rusty1s's avatar
update  
rusty1s committed
67
68
69

        N: int = 0
        if sparse_sizes is None or sparse_sizes[1] is None:
rusty1s's avatar
rusty1s committed
70
            if col.numel() > 0:
rusty1s's avatar
update  
rusty1s committed
71
                N = int(col.max()) + 1
rusty1s's avatar
rusty1s committed
72
        else:
rusty1s's avatar
update  
rusty1s committed
73
74
75
            _N = sparse_sizes[1]
            assert _N is not None
            N = _N
rusty1s's avatar
rusty1s committed
76
            if col.numel() > 0:
77
                assert trust_data or int(col.max()) < N
rusty1s's avatar
update  
rusty1s committed
78
79

        sparse_sizes = (M, N)
rusty1s's avatar
rusty1s committed
80

rusty1s's avatar
rusty1s committed
81
82
83
84
85
        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
86
            row = row.contiguous()
rusty1s's avatar
rusty1s committed
87

rusty1s's avatar
rusty1s committed
88
        if rowptr is not None:
rusty1s's avatar
rusty1s committed
89
            assert rowptr.dtype == torch.long
rusty1s's avatar
rusty1s committed
90
91
            assert rowptr.device == col.device
            assert rowptr.dim() == 1
rusty1s's avatar
rusty1s committed
92
            assert rowptr.numel() - 1 == sparse_sizes[0]
rusty1s's avatar
rusty1s committed
93
            rowptr = rowptr.contiguous()
rusty1s's avatar
rusty1s committed
94

rusty1s's avatar
rusty1s committed
95
96
97
        if value is not None:
            assert value.device == col.device
            assert value.size(0) == col.size(0)
rusty1s's avatar
rusty1s committed
98
            value = value.contiguous()
rusty1s's avatar
rusty1s committed
99
100
101
102
103

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

rusty1s's avatar
rusty1s committed
107
        if colptr is not None:
rusty1s's avatar
rusty1s committed
108
            assert colptr.dtype == torch.long
rusty1s's avatar
rusty1s committed
109
110
            assert colptr.device == col.device
            assert colptr.dim() == 1
rusty1s's avatar
rusty1s committed
111
            assert colptr.numel() - 1 == sparse_sizes[1]
rusty1s's avatar
rusty1s committed
112
            colptr = colptr.contiguous()
rusty1s's avatar
rusty1s committed
113
114
115
116
117

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

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

rusty1s's avatar
rusty1s committed
128
129
        if csc2csr is not None:
            assert csc2csr.dtype == torch.long
rusty1s's avatar
rusty1s committed
130
            assert csc2csr.device == col.device
rusty1s's avatar
rusty1s committed
131
            assert csc2csr.dim() == 1
rusty1s's avatar
rusty1s committed
132
            assert csc2csr.numel() == col.size(0)
rusty1s's avatar
rusty1s committed
133
            csc2csr = csc2csr.contiguous()
rusty1s's avatar
rusty1s committed
134

rusty1s's avatar
rusty1s committed
135
136
137
        self._row = row
        self._rowptr = rowptr
        self._col = col
rusty1s's avatar
rusty1s committed
138
        self._value = value
rusty1s's avatar
rusty1s committed
139
        self._sparse_sizes = tuple(sparse_sizes)
rusty1s's avatar
rusty1s committed
140
        self._rowcount = rowcount
rusty1s's avatar
rusty1s committed
141
        self._colptr = colptr
rusty1s's avatar
rusty1s committed
142
        self._colcount = colcount
rusty1s's avatar
rusty1s committed
143
144
        self._csr2csc = csr2csc
        self._csc2csr = csc2csr
rusty1s's avatar
rusty1s committed
145

rusty1s's avatar
rusty1s committed
146
        if not is_sorted:
rusty1s's avatar
rusty1s committed
147
            idx = self._col.new_zeros(self._col.numel() + 1)
rusty1s's avatar
rusty1s committed
148
149
150
            idx[1:] = self.row()
            idx[1:] *= self._sparse_sizes[1]
            idx[1:] += self._col
rusty1s's avatar
rusty1s committed
151
            if (idx[1:] < idx[:-1]).any():
rusty1s's avatar
rusty1s committed
152
                perm = idx[1:].argsort()
rusty1s's avatar
rusty1s committed
153
                self._row = self.row()[perm]
rusty1s's avatar
rusty1s committed
154
                self._col = self._col[perm]
rusty1s's avatar
rusty1s committed
155
156
                if value is not None:
                    self._value = value[perm]
rusty1s's avatar
rusty1s committed
157
158
159
                self._csr2csc = None
                self._csc2csr = None

rusty1s's avatar
rusty1s committed
160
161
    @classmethod
    def empty(self):
rusty1s's avatar
rusty1s committed
162
163
164
165
166
        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,
167
                             is_sorted=True, trust_data=True)
rusty1s's avatar
rusty1s committed
168

rusty1s's avatar
rusty1s committed
169
    def has_row(self) -> bool:
rusty1s's avatar
rusty1s committed
170
        return self._row is not None
rusty1s's avatar
rusty1s committed
171

rusty1s's avatar
rusty1s committed
172
    def row(self):
rusty1s's avatar
rusty1s committed
173
174
175
        row = self._row
        if row is not None:
            return row
rusty1s's avatar
rusty1s committed
176

rusty1s's avatar
rusty1s committed
177
178
        rowptr = self._rowptr
        if rowptr is not None:
rusty1s's avatar
rusty1s committed
179
            row = torch.ops.torch_sparse.ptr2ind(rowptr, self._col.numel())
rusty1s's avatar
rusty1s committed
180
181
182
183
184
185
            self._row = row
            return row

        raise ValueError

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

rusty1s's avatar
rusty1s committed
188
189
190
191
    def rowptr(self) -> torch.Tensor:
        rowptr = self._rowptr
        if rowptr is not None:
            return rowptr
rusty1s's avatar
rusty1s committed
192

rusty1s's avatar
rusty1s committed
193
194
        row = self._row
        if row is not None:
rusty1s's avatar
rusty1s committed
195
            rowptr = torch.ops.torch_sparse.ind2ptr(row, self._sparse_sizes[0])
rusty1s's avatar
rusty1s committed
196
197
198
199
200
201
            self._rowptr = rowptr
            return rowptr

        raise ValueError

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

rusty1s's avatar
rusty1s committed
204
    def has_value(self) -> bool:
rusty1s's avatar
rusty1s committed
205
        return self._value is not None
rusty1s's avatar
rusty1s committed
206

rusty1s's avatar
rusty1s committed
207
    def value(self) -> Optional[torch.Tensor]:
rusty1s's avatar
rusty1s committed
208
209
        return self._value

210
    def set_value_(self, value: Optional[torch.Tensor],
rusty1s's avatar
rusty1s committed
211
212
                   layout: Optional[str] = None):
        if value is not None:
rusty1s's avatar
rusty1s committed
213
            if get_layout(layout) == 'csc':
rusty1s's avatar
rusty1s committed
214
215
216
217
                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
218

rusty1s's avatar
rusty1s committed
219
220
        self._value = value
        return self
rusty1s's avatar
rusty1s committed
221

222
    def set_value(self, value: Optional[torch.Tensor],
rusty1s's avatar
rusty1s committed
223
224
                  layout: Optional[str] = None):
        if value is not None:
rusty1s's avatar
rusty1s committed
225
            if get_layout(layout) == 'csc':
rusty1s's avatar
rusty1s committed
226
227
228
229
                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
230

231
232
233
234
235
236
237
238
239
240
241
242
243
        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,
            trust_data=True)
rusty1s's avatar
rusty1s committed
244

rusty1s's avatar
rusty1s committed
245
    def sparse_sizes(self) -> Tuple[int, int]:
rusty1s's avatar
rusty1s committed
246
        return self._sparse_sizes
rusty1s's avatar
rusty1s committed
247

rusty1s's avatar
rusty1s committed
248
249
    def sparse_size(self, dim: int) -> int:
        return self._sparse_sizes[dim]
rusty1s's avatar
rusty1s committed
250

rusty1s's avatar
rusty1s committed
251
    def sparse_resize(self, sparse_sizes: Tuple[int, int]):
rusty1s's avatar
rusty1s committed
252
253
        assert len(sparse_sizes) == 2
        old_sparse_sizes, nnz = self._sparse_sizes, self._col.numel()
rusty1s's avatar
rusty1s committed
254

rusty1s's avatar
rusty1s committed
255
        diff_0 = sparse_sizes[0] - old_sparse_sizes[0]
rusty1s's avatar
rusty1s committed
256
257
258
259
260
261
        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
262
        elif diff_0 < 0:
rusty1s's avatar
rusty1s committed
263
264
265
266
267
            if rowptr is not None:
                rowptr = rowptr[:-diff_0]
            if rowcount is not None:
                rowcount = rowcount[:-diff_0]

rusty1s's avatar
rusty1s committed
268
        diff_1 = sparse_sizes[1] - old_sparse_sizes[1]
rusty1s's avatar
rusty1s committed
269
270
271
272
273
274
        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
275
        elif diff_1 < 0:
rusty1s's avatar
rusty1s committed
276
277
278
279
280
            if colptr is not None:
                colptr = colptr[:-diff_1]
            if colcount is not None:
                colcount = colcount[:-diff_1]

281
282
283
284
285
286
287
288
289
290
291
292
293
        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,
            trust_data=True)
rusty1s's avatar
rusty1s committed
294

rusty1s's avatar
rusty1s committed
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
    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()

312
        row = torch.div(idx, num_cols, rounding_mode='floor')
rusty1s's avatar
rusty1s committed
313
        col = idx % num_cols
rusty1s's avatar
rusty1s committed
314
        assert row.dtype == torch.long and col.dtype == torch.long
rusty1s's avatar
rusty1s committed
315
316
317
318

        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,
319
                             csc2csr=None, is_sorted=True, trust_data=True)
rusty1s's avatar
rusty1s committed
320

rusty1s's avatar
rusty1s committed
321
    def has_rowcount(self) -> bool:
rusty1s's avatar
rusty1s committed
322
323
        return self._rowcount is not None

rusty1s's avatar
rusty1s committed
324
325
326
327
328
329
    def rowcount(self) -> torch.Tensor:
        rowcount = self._rowcount
        if rowcount is not None:
            return rowcount

        rowptr = self.rowptr()
330
        rowcount = rowptr[1:] - rowptr[:-1]
rusty1s's avatar
rusty1s committed
331
332
        self._rowcount = rowcount
        return rowcount
rusty1s's avatar
rusty1s committed
333

rusty1s's avatar
rusty1s committed
334
    def has_colptr(self) -> bool:
rusty1s's avatar
rusty1s committed
335
        return self._colptr is not None
rusty1s's avatar
rusty1s committed
336

rusty1s's avatar
rusty1s committed
337
338
339
    def colptr(self) -> torch.Tensor:
        colptr = self._colptr
        if colptr is not None:
rusty1s's avatar
rusty1s committed
340
            return colptr
rusty1s's avatar
rusty1s committed
341

rusty1s's avatar
rusty1s committed
342
343
        csr2csc = self._csr2csc
        if csr2csc is not None:
rusty1s's avatar
rusty1s committed
344
345
            colptr = torch.ops.torch_sparse.ind2ptr(self._col[csr2csc],
                                                    self._sparse_sizes[1])
rusty1s's avatar
rusty1s committed
346
        else:
rusty1s's avatar
rusty1s committed
347
            colptr = self._col.new_zeros(self._sparse_sizes[1] + 1)
rusty1s's avatar
rusty1s committed
348
349
350
351
352
            torch.cumsum(self.colcount(), dim=0, out=colptr[1:])
        self._colptr = colptr
        return colptr

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

rusty1s's avatar
rusty1s committed
355
356
357
358
359
360
361
    def colcount(self) -> torch.Tensor:
        colcount = self._colcount
        if colcount is not None:
            return colcount

        colptr = self._colptr
        if colptr is not None:
362
            colcount = colptr[1:] - colptr[:-1]
rusty1s's avatar
rusty1s committed
363
        else:
364
365
            colcount = scatter_add(torch.ones_like(self._col), self._col,
                                   dim_size=self._sparse_sizes[1])
rusty1s's avatar
rusty1s committed
366
367
        self._colcount = colcount
        return colcount
rusty1s's avatar
rusty1s committed
368

rusty1s's avatar
rusty1s committed
369
    def has_csr2csc(self) -> bool:
rusty1s's avatar
rusty1s committed
370
        return self._csr2csc is not None
rusty1s's avatar
rusty1s committed
371

rusty1s's avatar
rusty1s committed
372
373
374
375
    def csr2csc(self) -> torch.Tensor:
        csr2csc = self._csr2csc
        if csr2csc is not None:
            return csr2csc
rusty1s's avatar
rusty1s committed
376

rusty1s's avatar
rusty1s committed
377
        idx = self._sparse_sizes[0] * self._col + self.row()
rusty1s's avatar
rusty1s committed
378
379
380
381
382
        csr2csc = idx.argsort()
        self._csr2csc = csr2csc
        return csr2csc

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

rusty1s's avatar
rusty1s committed
385
386
387
388
    def csc2csr(self) -> torch.Tensor:
        csc2csr = self._csc2csr
        if csc2csr is not None:
            return csc2csr
rusty1s's avatar
rusty1s committed
389

rusty1s's avatar
rusty1s committed
390
391
392
        csc2csr = self.csr2csc().argsort()
        self._csc2csr = csc2csr
        return csc2csr
rusty1s's avatar
rusty1s committed
393

rusty1s's avatar
rusty1s committed
394
395
    def is_coalesced(self) -> bool:
        idx = self._col.new_full((self._col.numel() + 1, ), -1)
rusty1s's avatar
rusty1s committed
396
        idx[1:] = self._sparse_sizes[1] * self.row() + self._col
rusty1s's avatar
rusty1s committed
397
398
399
400
        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
401
        idx[1:] = self._sparse_sizes[1] * self.row() + self._col
rusty1s's avatar
rusty1s committed
402
        mask = idx[1:] > idx[:-1]
rusty1s's avatar
rusty1s committed
403

rusty1s's avatar
rusty1s committed
404
        if mask.all():  # Skip if indices are already coalesced.
rusty1s's avatar
rusty1s committed
405
406
            return self

rusty1s's avatar
rusty1s committed
407
408
        row = self.row()[mask]
        col = self._col[mask]
rusty1s's avatar
rusty1s committed
409

rusty1s's avatar
rusty1s committed
410
411
        value = self._value
        if value is not None:
rusty1s's avatar
reset  
rusty1s committed
412
            ptr = mask.nonzero().flatten()
rusty1s's avatar
rusty1s committed
413
            ptr = torch.cat([ptr, ptr.new_full((1, ), value.size(0))])
rusty1s's avatar
rusty1s committed
414
            value = segment_csr(value, ptr, reduce=reduce)
rusty1s's avatar
rusty1s committed
415

416
417
418
        return SparseStorage(row=row, rowptr=None, col=col, value=value,
                             sparse_sizes=self._sparse_sizes, rowcount=None,
                             colptr=None, colcount=None, csr2csc=None,
419
                             csc2csr=None, is_sorted=True, trust_data=True)
rusty1s's avatar
rusty1s committed
420
421
422
423
424
425
426
427
428

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

rusty1s's avatar
rusty1s committed
431
432
433
434
435
436
    def clear_cache_(self):
        self._rowcount = None
        self._colptr = None
        self._colcount = None
        self._csr2csc = None
        self._csc2csr = None
rusty1s's avatar
rusty1s committed
437
        return self
rusty1s's avatar
rusty1s committed
438

rusty1s's avatar
rusty1s committed
439
440
    def cached_keys(self) -> List[str]:
        keys: List[str] = []
rusty1s's avatar
rusty1s committed
441
        if self.has_rowcount():
rusty1s's avatar
rusty1s committed
442
            keys.append('rowcount')
rusty1s's avatar
rusty1s committed
443
        if self.has_colptr():
rusty1s's avatar
rusty1s committed
444
            keys.append('colptr')
rusty1s's avatar
rusty1s committed
445
        if self.has_colcount():
rusty1s's avatar
rusty1s committed
446
            keys.append('colcount')
rusty1s's avatar
rusty1s committed
447
        if self.has_csr2csc():
rusty1s's avatar
rusty1s committed
448
            keys.append('csr2csc')
rusty1s's avatar
rusty1s committed
449
        if self.has_csc2csr():
rusty1s's avatar
rusty1s committed
450
451
452
453
454
            keys.append('csc2csr')
        return keys

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

rusty1s's avatar
rusty1s committed
456
    def copy(self):
457
458
459
460
461
462
463
464
465
466
467
468
469
        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,
            trust_data=True)
rusty1s's avatar
rusty1s committed
470

rusty1s's avatar
test  
rusty1s committed
471
    def clone(self):
rusty1s's avatar
rusty1s committed
472
473
474
475
476
477
        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
478
        col = self._col.clone()
rusty1s's avatar
rusty1s committed
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
        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
497

498
499
500
501
        return SparseStorage(row=row, rowptr=rowptr, col=col, value=value,
                             sparse_sizes=self._sparse_sizes,
                             rowcount=rowcount, colptr=colptr,
                             colcount=colcount, csr2csc=csr2csc,
502
                             csc2csr=csc2csr, is_sorted=True, trust_data=True)
rusty1s's avatar
rusty1s committed
503

504
    def type(self, dtype: torch.dtype, non_blocking: bool = False):
rusty1s's avatar
rusty1s committed
505
506
        value = self._value
        if value is not None:
507
            if dtype == value.dtype:
rusty1s's avatar
rusty1s committed
508
509
                return self
            else:
510
511
512
513
514
                return self.set_value(
                    value.to(
                        dtype=dtype,
                        non_blocking=non_blocking),
                    layout='coo')
rusty1s's avatar
rusty1s committed
515
516
517
        else:
            return self

518
519
520
521
522
    def type_as(self, tensor: torch.Tensor, non_blocking: bool = False):
        return self.type(dtype=tensor.dtype, non_blocking=non_blocking)

    def to_device(self, device: torch.device, non_blocking: bool = False):
        if device == self._col.device:
rusty1s's avatar
rusty1s committed
523
524
525
526
            return self

        row = self._row
        if row is not None:
527
            row = row.to(device, non_blocking=non_blocking)
rusty1s's avatar
rusty1s committed
528
529
        rowptr = self._rowptr
        if rowptr is not None:
530
531
            rowptr = rowptr.to(device, non_blocking=non_blocking)
        col = self._col.to(device, non_blocking=non_blocking)
rusty1s's avatar
rusty1s committed
532
533
        value = self._value
        if value is not None:
534
            value = value.to(device, non_blocking=non_blocking)
rusty1s's avatar
rusty1s committed
535
536
        rowcount = self._rowcount
        if rowcount is not None:
537
            rowcount = rowcount.to(device, non_blocking=non_blocking)
rusty1s's avatar
rusty1s committed
538
539
        colptr = self._colptr
        if colptr is not None:
540
            colptr = colptr.to(device, non_blocking=non_blocking)
rusty1s's avatar
rusty1s committed
541
542
        colcount = self._colcount
        if colcount is not None:
543
            colcount = colcount.to(device, non_blocking=non_blocking)
rusty1s's avatar
rusty1s committed
544
545
        csr2csc = self._csr2csc
        if csr2csc is not None:
546
            csr2csc = csr2csc.to(device, non_blocking=non_blocking)
rusty1s's avatar
rusty1s committed
547
548
        csc2csr = self._csc2csr
        if csc2csr is not None:
549
            csc2csr = csc2csr.to(device, non_blocking=non_blocking)
rusty1s's avatar
rusty1s committed
550

551
552
553
554
        return SparseStorage(row=row, rowptr=rowptr, col=col, value=value,
                             sparse_sizes=self._sparse_sizes,
                             rowcount=rowcount, colptr=colptr,
                             colcount=colcount, csr2csc=csr2csc,
555
556
557
558
                             csc2csr=csc2csr, is_sorted=True, trust_data=True)

    def device_as(self, tensor: torch.Tensor, non_blocking: bool = False):
        return self.to_device(device=tensor.device, non_blocking=non_blocking)
rusty1s's avatar
rusty1s committed
559

rusty1s's avatar
rusty1s committed
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
    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,
594
                             csc2csr=csc2csr, is_sorted=True, trust_data=True)
rusty1s's avatar
rusty1s committed
595

rusty1s's avatar
rusty1s committed
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
    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
622

623
624
625
626
        return SparseStorage(row=row, rowptr=rowptr, col=col, value=value,
                             sparse_sizes=self._sparse_sizes,
                             rowcount=rowcount, colptr=colptr,
                             colcount=colcount, csr2csc=csr2csc,
627
                             csc2csr=csc2csr, is_sorted=True, trust_data=True)
rusty1s's avatar
rusty1s committed
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
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718

    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