storage.py 22.7 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
54
55
56
57
            if rowptr is not None:
                M = rowptr.numel() - 1
            elif row is not None:
                M = row.max().item() + 1
            else:
                raise ValueError
rusty1s's avatar
rusty1s committed
58
            N = col.max().item() + 1
rusty1s's avatar
rusty1s committed
59
            sparse_sizes = (int(M), int(N))
rusty1s's avatar
rusty1s committed
60
        else:
rusty1s's avatar
rusty1s committed
61
            assert len(sparse_sizes) == 2
rusty1s's avatar
rusty1s committed
62

rusty1s's avatar
rusty1s committed
63
64
65
66
67
        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
68
            row = row.contiguous()
rusty1s's avatar
rusty1s committed
69

rusty1s's avatar
rusty1s committed
70
        if rowptr is not None:
rusty1s's avatar
rusty1s committed
71
            assert rowptr.dtype == torch.long
rusty1s's avatar
rusty1s committed
72
73
            assert rowptr.device == col.device
            assert rowptr.dim() == 1
rusty1s's avatar
rusty1s committed
74
            assert rowptr.numel() - 1 == sparse_sizes[0]
rusty1s's avatar
rusty1s committed
75
            rowptr = rowptr.contiguous()
rusty1s's avatar
rusty1s committed
76

rusty1s's avatar
rusty1s committed
77
78
79
        if value is not None:
            assert value.device == col.device
            assert value.size(0) == col.size(0)
rusty1s's avatar
rusty1s committed
80
            value = value.contiguous()
rusty1s's avatar
rusty1s committed
81
82
83
84
85

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

rusty1s's avatar
rusty1s committed
89
        if colptr is not None:
rusty1s's avatar
rusty1s committed
90
            assert colptr.dtype == torch.long
rusty1s's avatar
rusty1s committed
91
92
            assert colptr.device == col.device
            assert colptr.dim() == 1
rusty1s's avatar
rusty1s committed
93
            assert colptr.numel() - 1 == sparse_sizes[1]
rusty1s's avatar
rusty1s committed
94
            colptr = colptr.contiguous()
rusty1s's avatar
rusty1s committed
95
96
97
98
99

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

rusty1s's avatar
rusty1s committed
103
104
        if csr2csc is not None:
            assert csr2csc.dtype == torch.long
rusty1s's avatar
rusty1s committed
105
            assert csr2csc.device == col.device
rusty1s's avatar
rusty1s committed
106
            assert csr2csc.dim() == 1
rusty1s's avatar
rusty1s committed
107
            assert csr2csc.numel() == col.size(0)
rusty1s's avatar
rusty1s committed
108
            csr2csc = csr2csc.contiguous()
rusty1s's avatar
rusty1s committed
109

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

rusty1s's avatar
rusty1s committed
117
118
119
        self._row = row
        self._rowptr = rowptr
        self._col = col
rusty1s's avatar
rusty1s committed
120
        self._value = value
rusty1s's avatar
rusty1s committed
121
        self._sparse_sizes = tuple(sparse_sizes)
rusty1s's avatar
rusty1s committed
122
        self._rowcount = rowcount
rusty1s's avatar
rusty1s committed
123
        self._colptr = colptr
rusty1s's avatar
rusty1s committed
124
        self._colcount = colcount
rusty1s's avatar
rusty1s committed
125
126
        self._csr2csc = csr2csc
        self._csc2csr = csc2csr
rusty1s's avatar
rusty1s committed
127

rusty1s's avatar
rusty1s committed
128
        if not is_sorted:
rusty1s's avatar
rusty1s committed
129
            idx = self._col.new_zeros(self._col.numel() + 1)
rusty1s's avatar
rusty1s committed
130
131
132
            idx[1:] = self.row()
            idx[1:] *= self._sparse_sizes[1]
            idx[1:] += self._col
rusty1s's avatar
rusty1s committed
133
            if (idx[1:] < idx[:-1]).any():
rusty1s's avatar
rusty1s committed
134
                perm = idx[1:].argsort()
rusty1s's avatar
rusty1s committed
135
                self._row = self.row()[perm]
rusty1s's avatar
rusty1s committed
136
                self._col = self._col[perm]
rusty1s's avatar
rusty1s committed
137
138
                if value is not None:
                    self._value = value[perm]
rusty1s's avatar
rusty1s committed
139
140
141
                self._csr2csc = None
                self._csc2csr = None

rusty1s's avatar
rusty1s committed
142
143
144
145
146
147
148
149
150
151
152
153
154
    @classmethod
    def empty(self):
        self = SparseStorage.__new__(SparseStorage)
        self._row = None
        self._rowptr = None
        self._value = None
        self._rowcount = None
        self._colptr = None
        self._colcount = None
        self._csr2csc = None
        self._csc2csr = None
        return self

rusty1s's avatar
rusty1s committed
155
    def has_row(self) -> bool:
rusty1s's avatar
rusty1s committed
156
        return self._row is not None
rusty1s's avatar
rusty1s committed
157

rusty1s's avatar
rusty1s committed
158
    def row(self):
rusty1s's avatar
rusty1s committed
159
160
161
        row = self._row
        if row is not None:
            return row
rusty1s's avatar
rusty1s committed
162

rusty1s's avatar
rusty1s committed
163
164
        rowptr = self._rowptr
        if rowptr is not None:
rusty1s's avatar
rusty1s committed
165
            row = torch.ops.torch_sparse.ptr2ind(rowptr, self._col.numel())
rusty1s's avatar
rusty1s committed
166
167
168
169
170
171
            self._row = row
            return row

        raise ValueError

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

rusty1s's avatar
rusty1s committed
174
175
176
177
    def rowptr(self) -> torch.Tensor:
        rowptr = self._rowptr
        if rowptr is not None:
            return rowptr
rusty1s's avatar
rusty1s committed
178

rusty1s's avatar
rusty1s committed
179
180
        row = self._row
        if row is not None:
rusty1s's avatar
rusty1s committed
181
            rowptr = torch.ops.torch_sparse.ind2ptr(row, self._sparse_sizes[0])
rusty1s's avatar
rusty1s committed
182
183
184
185
186
187
            self._rowptr = rowptr
            return rowptr

        raise ValueError

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

rusty1s's avatar
rusty1s committed
190
    def has_value(self) -> bool:
rusty1s's avatar
rusty1s committed
191
        return self._value is not None
rusty1s's avatar
rusty1s committed
192

rusty1s's avatar
rusty1s committed
193
    def value(self) -> Optional[torch.Tensor]:
rusty1s's avatar
rusty1s committed
194
195
        return self._value

196
    def set_value_(self, value: Optional[torch.Tensor],
rusty1s's avatar
rusty1s committed
197
198
                   layout: Optional[str] = None):
        if value is not None:
rusty1s's avatar
rusty1s committed
199
            if get_layout(layout) == 'csc':
rusty1s's avatar
rusty1s committed
200
201
202
203
                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
204

rusty1s's avatar
rusty1s committed
205
206
        self._value = value
        return self
rusty1s's avatar
rusty1s committed
207

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

217
218
219
220
221
        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
222

rusty1s's avatar
rusty1s committed
223
    def sparse_sizes(self) -> Tuple[int, int]:
rusty1s's avatar
rusty1s committed
224
        return self._sparse_sizes
rusty1s's avatar
rusty1s committed
225

rusty1s's avatar
rusty1s committed
226
227
    def sparse_size(self, dim: int) -> int:
        return self._sparse_sizes[dim]
rusty1s's avatar
rusty1s committed
228

rusty1s's avatar
rusty1s committed
229
    def sparse_resize(self, sparse_sizes: Tuple[int, int]):
rusty1s's avatar
rusty1s committed
230
231
        assert len(sparse_sizes) == 2
        old_sparse_sizes, nnz = self._sparse_sizes, self._col.numel()
rusty1s's avatar
rusty1s committed
232

rusty1s's avatar
rusty1s committed
233
        diff_0 = sparse_sizes[0] - old_sparse_sizes[0]
rusty1s's avatar
rusty1s committed
234
235
236
237
238
239
        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
240
        elif diff_0 < 0:
rusty1s's avatar
rusty1s committed
241
242
243
244
245
            if rowptr is not None:
                rowptr = rowptr[:-diff_0]
            if rowcount is not None:
                rowcount = rowcount[:-diff_0]

rusty1s's avatar
rusty1s committed
246
        diff_1 = sparse_sizes[1] - old_sparse_sizes[1]
rusty1s's avatar
rusty1s committed
247
248
249
250
251
252
        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
253
        elif diff_1 < 0:
rusty1s's avatar
rusty1s committed
254
255
256
257
258
            if colptr is not None:
                colptr = colptr[:-diff_1]
            if colcount is not None:
                colcount = colcount[:-diff_1]

259
260
261
262
263
        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
264

rusty1s's avatar
rusty1s committed
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
    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
282
        row = idx // num_cols
rusty1s's avatar
rusty1s committed
283
        col = idx % num_cols
rusty1s's avatar
rusty1s committed
284
        assert row.dtype == torch.long and col.dtype == torch.long
rusty1s's avatar
rusty1s committed
285
286
287
288
289
290

        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
291
    def has_rowcount(self) -> bool:
rusty1s's avatar
rusty1s committed
292
293
        return self._rowcount is not None

rusty1s's avatar
rusty1s committed
294
295
296
297
298
299
    def rowcount(self) -> torch.Tensor:
        rowcount = self._rowcount
        if rowcount is not None:
            return rowcount

        rowptr = self.rowptr()
300
        rowcount = rowptr[1:] - rowptr[:-1]
rusty1s's avatar
rusty1s committed
301
302
        self._rowcount = rowcount
        return rowcount
rusty1s's avatar
rusty1s committed
303

rusty1s's avatar
rusty1s committed
304
    def has_colptr(self) -> bool:
rusty1s's avatar
rusty1s committed
305
        return self._colptr is not None
rusty1s's avatar
rusty1s committed
306

rusty1s's avatar
rusty1s committed
307
308
309
    def colptr(self) -> torch.Tensor:
        colptr = self._colptr
        if colptr is not None:
rusty1s's avatar
rusty1s committed
310
            return colptr
rusty1s's avatar
rusty1s committed
311

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

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

rusty1s's avatar
rusty1s committed
325
326
327
328
329
330
331
    def colcount(self) -> torch.Tensor:
        colcount = self._colcount
        if colcount is not None:
            return colcount

        colptr = self._colptr
        if colptr is not None:
332
            colcount = colptr[1:] - colptr[:-1]
rusty1s's avatar
rusty1s committed
333
        else:
334
335
            colcount = scatter_add(torch.ones_like(self._col), self._col,
                                   dim_size=self._sparse_sizes[1])
rusty1s's avatar
rusty1s committed
336
337
        self._colcount = colcount
        return colcount
rusty1s's avatar
rusty1s committed
338

rusty1s's avatar
rusty1s committed
339
    def has_csr2csc(self) -> bool:
rusty1s's avatar
rusty1s committed
340
        return self._csr2csc is not None
rusty1s's avatar
rusty1s committed
341

rusty1s's avatar
rusty1s committed
342
343
344
345
    def csr2csc(self) -> torch.Tensor:
        csr2csc = self._csr2csc
        if csr2csc is not None:
            return csr2csc
rusty1s's avatar
rusty1s committed
346

rusty1s's avatar
rusty1s committed
347
        idx = self._sparse_sizes[0] * self._col + self.row()
rusty1s's avatar
rusty1s committed
348
349
350
351
352
        csr2csc = idx.argsort()
        self._csr2csc = csr2csc
        return csr2csc

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

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

rusty1s's avatar
rusty1s committed
360
361
362
        csc2csr = self.csr2csc().argsort()
        self._csc2csr = csc2csr
        return csc2csr
rusty1s's avatar
rusty1s committed
363

rusty1s's avatar
rusty1s committed
364
365
    def is_coalesced(self) -> bool:
        idx = self._col.new_full((self._col.numel() + 1, ), -1)
rusty1s's avatar
rusty1s committed
366
        idx[1:] = self._sparse_sizes[1] * self.row() + self._col
rusty1s's avatar
rusty1s committed
367
368
369
370
        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
371
        idx[1:] = self._sparse_sizes[1] * self.row() + self._col
rusty1s's avatar
rusty1s committed
372
        mask = idx[1:] > idx[:-1]
rusty1s's avatar
rusty1s committed
373

rusty1s's avatar
rusty1s committed
374
        if mask.all():  # Skip if indices are already coalesced.
rusty1s's avatar
rusty1s committed
375
376
            return self

rusty1s's avatar
rusty1s committed
377
378
        row = self.row()[mask]
        col = self._col[mask]
rusty1s's avatar
rusty1s committed
379

rusty1s's avatar
rusty1s committed
380
381
        value = self._value
        if value is not None:
rusty1s's avatar
rusty1s committed
382
            ptr = mask.nonzero(as_tuple=False).flatten()
rusty1s's avatar
rusty1s committed
383
            ptr = torch.cat([ptr, ptr.new_full((1, ), value.size(0))])
rusty1s's avatar
rusty1s committed
384
            value = segment_csr(value, ptr, reduce=reduce)
rusty1s's avatar
rusty1s committed
385
386
            value = value[0] if isinstance(value, tuple) else value

387
388
389
390
        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
391
392
393
394
395
396
397
398
399

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

rusty1s's avatar
rusty1s committed
402
403
404
405
406
407
    def clear_cache_(self):
        self._rowcount = None
        self._colptr = None
        self._colcount = None
        self._csr2csc = None
        self._csc2csr = None
rusty1s's avatar
rusty1s committed
408
        return self
rusty1s's avatar
rusty1s committed
409

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

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

rusty1s's avatar
rusty1s committed
427
    def copy(self):
428
429
430
431
432
433
        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
434

rusty1s's avatar
test  
rusty1s committed
435
    def clone(self):
rusty1s's avatar
rusty1s committed
436
437
438
439
440
441
        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
442
        col = self._col.clone()
rusty1s's avatar
rusty1s committed
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
        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
461

462
463
464
465
466
        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
467

rusty1s's avatar
rusty1s committed
468
    def type_as(self, tensor: torch.Tensor):
rusty1s's avatar
rusty1s committed
469
470
471
472
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
        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
507

508
509
510
511
512
        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
513

rusty1s's avatar
rusty1s committed
514
515
516
517
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
    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
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
    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
576

577
578
579
580
581
        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
582
583
584
585
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

    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