tensor.py 20.6 KB
Newer Older
rusty1s's avatar
repr  
rusty1s committed
1
from textwrap import indent
rusty1s's avatar
typing  
rusty1s committed
2
from typing import Optional, List, Tuple, Dict, Union, Any
rusty1s's avatar
rusty1s committed
3
4
5
6

import torch
import scipy.sparse

rusty1s's avatar
rusty1s committed
7
from torch_sparse.storage import SparseStorage, get_layout
rusty1s's avatar
rusty1s committed
8
from torch_sparse.utils import is_scalar
rusty1s's avatar
rusty1s committed
9
10


rusty1s's avatar
rusty1s committed
11
@torch.jit.script
rusty1s's avatar
rusty1s committed
12
class SparseTensor(object):
rusty1s's avatar
rusty1s committed
13
14
    storage: SparseStorage

rusty1s's avatar
rusty1s committed
15
    def __init__(self, row: Optional[torch.Tensor] = None,
rusty1s's avatar
rusty1s committed
16
17
18
                 rowptr: Optional[torch.Tensor] = None,
                 col: Optional[torch.Tensor] = None,
                 value: Optional[torch.Tensor] = None,
rusty1s's avatar
rusty1s committed
19
20
                 sparse_sizes: Optional[Tuple[int, int]] = None,
                 is_sorted: bool = False):
rusty1s's avatar
rusty1s committed
21
22
23
24
25
        self.storage = SparseStorage(row=row, rowptr=rowptr, col=col,
                                     value=value, sparse_sizes=sparse_sizes,
                                     rowcount=None, colptr=None, colcount=None,
                                     csr2csc=None, csc2csr=None,
                                     is_sorted=is_sorted)
rusty1s's avatar
rusty1s committed
26
27

    @classmethod
rusty1s's avatar
rusty1s committed
28
    def from_storage(self, storage: SparseStorage):
rusty1s's avatar
rusty1s committed
29
        self = SparseTensor.__new__(SparseTensor)
rusty1s's avatar
rusty1s committed
30
        self.storage = storage
rusty1s's avatar
rusty1s committed
31
32
        return self

rusty1s's avatar
rusty1s committed
33
34
35
36
37
38
39
40
41
    @classmethod
    def from_edge_index(self, edge_index: torch.Tensor,
                        edge_attr: Optional[torch.Tensor] = None,
                        sparse_sizes: Optional[Tuple[int, int]] = None,
                        is_sorted: bool = False):
        return SparseTensor(row=edge_index[0], rowptr=None, col=edge_index[1],
                            value=edge_attr, sparse_sizes=sparse_sizes,
                            is_sorted=is_sorted)

rusty1s's avatar
rusty1s committed
42
    @classmethod
rusty1s's avatar
rusty1s committed
43
    def from_dense(self, mat: torch.Tensor, has_value: bool = True):
rusty1s's avatar
rusty1s committed
44
45
46
47
        if mat.dim() > 2:
            index = mat.abs().sum([i for i in range(2, mat.dim())]).nonzero()
        else:
            index = mat.nonzero()
rusty1s's avatar
rusty1s committed
48
        index = index.t()
rusty1s's avatar
rusty1s committed
49

rusty1s's avatar
rusty1s committed
50
51
52
53
54
55
56
        row = index[0]
        col = index[1]

        value: Optional[torch.Tensor] = None
        if has_value:
            value = mat[row, col]

rusty1s's avatar
rusty1s committed
57
58
59
        return SparseTensor(row=row, rowptr=None, col=col, value=value,
                            sparse_sizes=(mat.size(0), mat.size(1)),
                            is_sorted=True)
rusty1s's avatar
rusty1s committed
60
61

    @classmethod
rusty1s's avatar
rusty1s committed
62
    def from_torch_sparse_coo_tensor(self, mat: torch.Tensor,
rusty1s's avatar
rusty1s committed
63
                                     has_value: bool = True):
rusty1s's avatar
rusty1s committed
64
65
66
        mat = mat.coalesce()
        index = mat._indices()
        row, col = index[0], index[1]
rusty1s's avatar
rusty1s committed
67
68
69
70
71

        value: Optional[torch.Tensor] = None
        if has_value:
            value = mat._values()

rusty1s's avatar
rusty1s committed
72
73
74
        return SparseTensor(row=row, rowptr=None, col=col, value=value,
                            sparse_sizes=(mat.size(0), mat.size(1)),
                            is_sorted=True)
rusty1s's avatar
rusty1s committed
75
76

    @classmethod
rusty1s's avatar
rusty1s committed
77
78
    def eye(self, M: int, N: Optional[int] = None,
            options: Optional[torch.Tensor] = None, has_value: bool = True,
rusty1s's avatar
rusty1s committed
79
            fill_cache: bool = False):
rusty1s's avatar
rusty1s committed
80

rusty1s's avatar
rusty1s committed
81
        N = M if N is None else N
rusty1s's avatar
rusty1s committed
82

rusty1s's avatar
rusty1s committed
83
84
85
86
        if options is not None:
            row = torch.arange(min(M, N), device=options.device)
        else:
            row = torch.arange(min(M, N))
rusty1s's avatar
rusty1s committed
87
        col = row
rusty1s's avatar
rusty1s committed
88

rusty1s's avatar
rusty1s committed
89
90
        rowptr = torch.arange(M + 1, dtype=torch.long, device=row.device)
        if M > N:
rusty1s's avatar
rusty1s committed
91
            rowptr[N + 1:] = N
rusty1s's avatar
rusty1s committed
92
93

        value: Optional[torch.Tensor] = None
rusty1s's avatar
rusty1s committed
94
        if has_value:
rusty1s's avatar
rusty1s committed
95
            if options is not None:
rusty1s's avatar
rusty1s committed
96
97
                value = torch.ones(row.numel(), dtype=options.dtype,
                                   device=row.device)
rusty1s's avatar
rusty1s committed
98
99
100
101
102
103
104
105
            else:
                value = torch.ones(row.numel(), device=row.device)

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

        if fill_cache:
rusty1s's avatar
rusty1s committed
108
            rowcount = torch.ones(M, dtype=torch.long, device=row.device)
rusty1s's avatar
rusty1s committed
109
            if M > N:
rusty1s's avatar
rusty1s committed
110
111
112
113
                rowcount[N:] = 0

            colptr = torch.arange(N + 1, dtype=torch.long, device=row.device)
            colcount = torch.ones(N, dtype=torch.long, device=row.device)
rusty1s's avatar
rusty1s committed
114
            if N > M:
rusty1s's avatar
rusty1s committed
115
116
                colptr[M + 1:] = M
                colcount[M:] = 0
rusty1s's avatar
rusty1s committed
117
118
            csr2csc = csc2csr = row

rusty1s's avatar
rusty1s committed
119
        storage: SparseStorage = SparseStorage(
rusty1s's avatar
rusty1s committed
120
121
122
            row=row, rowptr=rowptr, col=col, value=value, sparse_sizes=(M, N),
            rowcount=rowcount, colptr=colptr, colcount=colcount,
            csr2csc=csr2csc, csc2csr=csc2csr, is_sorted=True)
rusty1s's avatar
rusty1s committed
123

rusty1s's avatar
rusty1s committed
124
125
126
127
128
        self = SparseTensor.__new__(SparseTensor)
        self.storage = storage
        return self

    def copy(self):
rusty1s's avatar
rusty1s committed
129
        return self.from_storage(self.storage)
rusty1s's avatar
rusty1s committed
130
131

    def clone(self):
rusty1s's avatar
rusty1s committed
132
        return self.from_storage(self.storage.clone())
rusty1s's avatar
rusty1s committed
133

rusty1s's avatar
rusty1s committed
134
135
136
137
138
139
140
141
142
143
    def type_as(self, tensor=torch.Tensor):
        value = self.storage._value
        if value is None or tensor.dtype == value.dtype:
            return self
        return self.from_storage(self.storage.type_as(tensor))

    def device_as(self, tensor: torch.Tensor, non_blocking: bool = False):
        if tensor.device == self.device():
            return self
        return self.from_storage(self.storage.device_as(tensor, non_blocking))
rusty1s's avatar
rusty1s committed
144
145
146

    # Formats #################################################################

rusty1s's avatar
rusty1s committed
147
148
    def coo(self) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
        return self.storage.row(), self.storage.col(), self.storage.value()
rusty1s's avatar
rusty1s committed
149

rusty1s's avatar
rusty1s committed
150
151
    def csr(self) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
        return self.storage.rowptr(), self.storage.col(), self.storage.value()
rusty1s's avatar
rusty1s committed
152

rusty1s's avatar
rusty1s committed
153
154
155
156
157
158
    def csc(self) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
        perm = self.storage.csr2csc()
        value = self.storage.value()
        if value is not None:
            value = value[perm]
        return self.storage.colptr(), self.storage.row()[perm], value
rusty1s's avatar
rusty1s committed
159
160
161

    # Storage inheritance #####################################################

rusty1s's avatar
rusty1s committed
162
    def has_value(self) -> bool:
rusty1s's avatar
rusty1s committed
163
        return self.storage.has_value()
rusty1s's avatar
rusty1s committed
164

rusty1s's avatar
rusty1s committed
165
    def set_value_(self, value: Optional[torch.Tensor],
rusty1s's avatar
rusty1s committed
166
167
                   layout: Optional[str] = None):
        self.storage.set_value_(value, layout)
rusty1s's avatar
rusty1s committed
168
169
        return self

rusty1s's avatar
rusty1s committed
170
    def set_value(self, value: Optional[torch.Tensor],
rusty1s's avatar
rusty1s committed
171
172
173
                  layout: Optional[str] = None):
        return self.from_storage(self.storage.set_value(value, layout))

rusty1s's avatar
rusty1s committed
174
    def sparse_sizes(self) -> Tuple[int, int]:
rusty1s's avatar
rusty1s committed
175
        return self.storage.sparse_sizes()
rusty1s's avatar
rusty1s committed
176

rusty1s's avatar
rusty1s committed
177
178
    def sparse_size(self, dim: int) -> int:
        return self.storage.sparse_sizes()[dim]
rusty1s's avatar
rusty1s committed
179

rusty1s's avatar
rusty1s committed
180
    def sparse_resize(self, sparse_sizes: Tuple[int, int]):
rusty1s's avatar
rusty1s committed
181
        return self.from_storage(self.storage.sparse_resize(sparse_sizes))
rusty1s's avatar
rusty1s committed
182

rusty1s's avatar
rusty1s committed
183
184
185
186
    def sparse_reshape(self, num_rows: int, num_cols: int):
        return self.from_storage(
            self.storage.sparse_reshape(num_rows, num_cols))

rusty1s's avatar
rusty1s committed
187
    def is_coalesced(self) -> bool:
rusty1s's avatar
rusty1s committed
188
        return self.storage.is_coalesced()
rusty1s's avatar
rusty1s committed
189

rusty1s's avatar
rusty1s committed
190
    def coalesce(self, reduce: str = "sum"):
rusty1s's avatar
rusty1s committed
191
        return self.from_storage(self.storage.coalesce(reduce))
rusty1s's avatar
rusty1s committed
192

rusty1s's avatar
rusty1s committed
193
194
    def fill_cache_(self):
        self.storage.fill_cache_()
rusty1s's avatar
rusty1s committed
195
196
        return self

rusty1s's avatar
rusty1s committed
197
198
    def clear_cache_(self):
        self.storage.clear_cache_()
rusty1s's avatar
rusty1s committed
199
200
201
202
        return self

    # Utility functions #######################################################

rusty1s's avatar
rusty1s committed
203
    def fill_value_(self, fill_value: float,
rusty1s's avatar
rusty1s committed
204
205
                    options: Optional[torch.Tensor] = None):
        if options is not None:
rusty1s's avatar
rusty1s committed
206
            value = torch.full((self.nnz(), ), fill_value, dtype=options.dtype,
rusty1s's avatar
rusty1s committed
207
208
                               device=self.device())
        else:
rusty1s's avatar
rusty1s committed
209
            value = torch.full((self.nnz(), ), fill_value,
rusty1s's avatar
rusty1s committed
210
211
212
                               device=self.device())
        return self.set_value_(value, layout='coo')

rusty1s's avatar
rusty1s committed
213
    def fill_value(self, fill_value: float,
rusty1s's avatar
rusty1s committed
214
215
                   options: Optional[torch.Tensor] = None):
        if options is not None:
rusty1s's avatar
rusty1s committed
216
            value = torch.full((self.nnz(), ), fill_value, dtype=options.dtype,
rusty1s's avatar
rusty1s committed
217
218
                               device=self.device())
        else:
rusty1s's avatar
rusty1s committed
219
            value = torch.full((self.nnz(), ), fill_value,
rusty1s's avatar
rusty1s committed
220
221
222
223
                               device=self.device())
        return self.set_value(value, layout='coo')

    def sizes(self) -> List[int]:
rusty1s's avatar
rusty1s committed
224
        sparse_sizes = self.sparse_sizes()
rusty1s's avatar
rusty1s committed
225
226
        value = self.storage.value()
        if value is not None:
rusty1s's avatar
rusty1s committed
227
228
229
            return list(sparse_sizes) + list(value.size())[1:]
        else:
            return list(sparse_sizes)
rusty1s's avatar
rusty1s committed
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245

    def size(self, dim: int) -> int:
        return self.sizes()[dim]

    def dim(self) -> int:
        return len(self.sizes())

    def nnz(self) -> int:
        return self.storage.col().numel()

    def numel(self) -> int:
        value = self.storage.value()
        if value is not None:
            return value.numel()
        else:
            return self.nnz()
rusty1s's avatar
rusty1s committed
246

rusty1s's avatar
rusty1s committed
247
    def density(self) -> float:
rusty1s's avatar
rusty1s committed
248
249
        return self.nnz() / (self.sparse_size(0) * self.sparse_size(1))

rusty1s's avatar
rusty1s committed
250
    def sparsity(self) -> float:
rusty1s's avatar
rusty1s committed
251
252
        return 1 - self.density()

rusty1s's avatar
rusty1s committed
253
    def avg_row_length(self) -> float:
rusty1s's avatar
rusty1s committed
254
255
        return self.nnz() / self.sparse_size(0)

rusty1s's avatar
rusty1s committed
256
    def avg_col_length(self) -> float:
rusty1s's avatar
rusty1s committed
257
258
        return self.nnz() / self.sparse_size(1)

rusty1s's avatar
rusty1s committed
259
260
261
262
    def bandwidth(self) -> int:
        row, col, _ = self.coo()
        return int((row - col).abs_().max())

rusty1s's avatar
rusty1s committed
263
264
265
266
    def avg_bandwidth(self) -> float:
        row, col, _ = self.coo()
        return float((row - col).abs_().to(torch.float).mean())

rusty1s's avatar
rusty1s committed
267
268
269
270
271
    def bandwidth_proportion(self, bandwidth: int) -> float:
        row, col, _ = self.coo()
        tmp = (row - col).abs_()
        return int((tmp <= bandwidth).sum()) / self.nnz()

rusty1s's avatar
rusty1s committed
272
    def is_quadratic(self) -> bool:
rusty1s's avatar
rusty1s committed
273
274
        return self.sparse_size(0) == self.sparse_size(1)

rusty1s's avatar
rusty1s committed
275
276
    def is_symmetric(self) -> bool:
        if not self.is_quadratic():
rusty1s's avatar
rusty1s committed
277
278
            return False

rusty1s's avatar
rusty1s committed
279
280
281
282
283
284
        rowptr, col, value1 = self.csr()
        colptr, row, value2 = self.csc()

        if (rowptr != colptr).any() or (col != row).any():
            return False

rusty1s's avatar
rusty1s committed
285
        if value1 is None or value2 is None:
rusty1s's avatar
rusty1s committed
286
            return True
rusty1s's avatar
rusty1s committed
287
288
        else:
            return bool((value1 == value2).all())
rusty1s's avatar
rusty1s committed
289

rusty1s's avatar
rusty1s committed
290
291
292
293
294
295
296
297
298
    def to_symmetric(self, reduce: str = "sum"):
        row, col, value = self.coo()

        row, col = torch.cat([row, col], dim=0), torch.cat([col, row], dim=0)
        if value is not None:
            value = torch.cat([value, value], dim=0)

        N = max(self.size(0), self.size(1))

rusty1s's avatar
rusty1s committed
299
300
        out = SparseTensor(row=row, rowptr=None, col=col, value=value,
                           sparse_sizes=(N, N), is_sorted=False)
rusty1s's avatar
rusty1s committed
301
302
303
        out = out.coalesce(reduce)
        return out

rusty1s's avatar
rusty1s committed
304
    def detach_(self):
rusty1s's avatar
rusty1s committed
305
306
307
        value = self.storage.value()
        if value is not None:
            value.detach_()
rusty1s's avatar
rusty1s committed
308
309
310
        return self

    def detach(self):
rusty1s's avatar
rusty1s committed
311
312
313
314
315
316
317
318
319
320
321
        value = self.storage.value()
        if value is not None:
            value = value.detach()
        return self.set_value(value, layout='coo')

    def requires_grad(self) -> bool:
        value = self.storage.value()
        if value is not None:
            return value.requires_grad
        else:
            return False
rusty1s's avatar
rusty1s committed
322

rusty1s's avatar
rusty1s committed
323
    def requires_grad_(self, requires_grad: bool = True,
rusty1s's avatar
rusty1s committed
324
                       options: Optional[torch.Tensor] = None):
rusty1s's avatar
rusty1s committed
325
        if requires_grad and not self.has_value():
rusty1s's avatar
rusty1s committed
326
            self.fill_value_(1., options=options)
rusty1s's avatar
rusty1s committed
327

rusty1s's avatar
rusty1s committed
328
329
330
        value = self.storage.value()
        if value is not None:
            value.requires_grad_(requires_grad)
rusty1s's avatar
rusty1s committed
331
332
        return self

rusty1s's avatar
rusty1s committed
333
    def pin_memory(self):
rusty1s's avatar
rusty1s committed
334
        return self.from_storage(self.storage.pin_memory())
rusty1s's avatar
rusty1s committed
335

rusty1s's avatar
rusty1s committed
336
337
    def is_pinned(self) -> bool:
        return self.storage.is_pinned()
rusty1s's avatar
rusty1s committed
338

rusty1s's avatar
rusty1s committed
339
340
341
342
343
    def options(self) -> torch.Tensor:
        value = self.storage.value()
        if value is not None:
            return value
        else:
rusty1s's avatar
rusty1s committed
344
345
            return torch.tensor(0., dtype=torch.float,
                                device=self.storage.col().device)
rusty1s's avatar
rusty1s committed
346
347

    def device(self):
rusty1s's avatar
rusty1s committed
348
        return self.storage.col().device
rusty1s's avatar
rusty1s committed
349
350

    def cpu(self):
rusty1s's avatar
rusty1s committed
351
        return self.device_as(torch.tensor(0.), non_blocking=False)
rusty1s's avatar
rusty1s committed
352

rusty1s's avatar
rusty1s committed
353
    def cuda(self, options: Optional[torch.Tensor] = None,
rusty1s's avatar
rusty1s committed
354
             non_blocking: bool = False):
rusty1s's avatar
rusty1s committed
355
356
        if options is not None:
            return self.device_as(options, non_blocking)
rusty1s's avatar
rusty1s committed
357
        else:
rusty1s's avatar
rusty1s committed
358
359
            options = torch.tensor(0.).cuda()
            return self.device_as(options, non_blocking)
rusty1s's avatar
rusty1s committed
360

rusty1s's avatar
rusty1s committed
361
362
    def is_cuda(self) -> bool:
        return self.storage.col().is_cuda
rusty1s's avatar
rusty1s committed
363

rusty1s's avatar
rusty1s committed
364
365
    def dtype(self):
        return self.options().dtype
rusty1s's avatar
rusty1s committed
366

rusty1s's avatar
rusty1s committed
367
368
    def is_floating_point(self) -> bool:
        return torch.is_floating_point(self.options())
rusty1s's avatar
rusty1s committed
369
370

    def bfloat16(self):
rusty1s's avatar
rusty1s committed
371
372
        return self.type_as(
            torch.tensor(0, dtype=torch.bfloat16, device=self.device()))
rusty1s's avatar
rusty1s committed
373
374

    def bool(self):
rusty1s's avatar
rusty1s committed
375
376
        return self.type_as(
            torch.tensor(0, dtype=torch.bool, device=self.device()))
rusty1s's avatar
rusty1s committed
377
378

    def byte(self):
rusty1s's avatar
rusty1s committed
379
380
        return self.type_as(
            torch.tensor(0, dtype=torch.uint8, device=self.device()))
rusty1s's avatar
rusty1s committed
381
382

    def char(self):
rusty1s's avatar
rusty1s committed
383
384
        return self.type_as(
            torch.tensor(0, dtype=torch.int8, device=self.device()))
rusty1s's avatar
rusty1s committed
385
386

    def half(self):
rusty1s's avatar
rusty1s committed
387
388
        return self.type_as(
            torch.tensor(0, dtype=torch.half, device=self.device()))
rusty1s's avatar
rusty1s committed
389
390

    def float(self):
rusty1s's avatar
rusty1s committed
391
392
        return self.type_as(
            torch.tensor(0, dtype=torch.float, device=self.device()))
rusty1s's avatar
rusty1s committed
393
394

    def double(self):
rusty1s's avatar
rusty1s committed
395
396
        return self.type_as(
            torch.tensor(0, dtype=torch.double, device=self.device()))
rusty1s's avatar
rusty1s committed
397
398

    def short(self):
rusty1s's avatar
rusty1s committed
399
400
        return self.type_as(
            torch.tensor(0, dtype=torch.short, device=self.device()))
rusty1s's avatar
rusty1s committed
401
402

    def int(self):
rusty1s's avatar
rusty1s committed
403
404
        return self.type_as(
            torch.tensor(0, dtype=torch.int, device=self.device()))
rusty1s's avatar
rusty1s committed
405
406

    def long(self):
rusty1s's avatar
rusty1s committed
407
408
        return self.type_as(
            torch.tensor(0, dtype=torch.long, device=self.device()))
rusty1s's avatar
rusty1s committed
409
410
411

    # Conversions #############################################################

rusty1s's avatar
rusty1s committed
412
    def to_dense(self, options: Optional[torch.Tensor] = None):
rusty1s's avatar
rusty1s committed
413
        row, col, value = self.coo()
rusty1s's avatar
rusty1s committed
414

rusty1s's avatar
fixes  
rusty1s committed
415
        if value is not None:
rusty1s's avatar
rusty1s committed
416
417
            mat = torch.zeros(self.sizes(), dtype=value.dtype,
                              device=self.device())
rusty1s's avatar
fixes  
rusty1s committed
418
        elif options is not None:
rusty1s's avatar
rusty1s committed
419
420
            mat = torch.zeros(self.sizes(), dtype=options.dtype,
                              device=self.device())
rusty1s's avatar
rusty1s committed
421
422
423
424
425
426
        else:
            mat = torch.zeros(self.sizes(), device=self.device())

        if value is not None:
            mat[row, col] = value
        else:
rusty1s's avatar
rusty1s committed
427
428
            mat[row, col] = torch.ones(self.nnz(), dtype=mat.dtype,
                                       device=mat.device)
rusty1s's avatar
rusty1s committed
429

rusty1s's avatar
rusty1s committed
430
431
        return mat

rusty1s's avatar
rusty1s committed
432
433
    def to_torch_sparse_coo_tensor(self,
                                   options: Optional[torch.Tensor] = None):
rusty1s's avatar
rusty1s committed
434
435
436
        row, col, value = self.coo()
        index = torch.stack([row, col], dim=0)
        if value is None:
rusty1s's avatar
rusty1s committed
437
            if options is not None:
rusty1s's avatar
rusty1s committed
438
439
                value = torch.ones(self.nnz(), dtype=options.dtype,
                                   device=self.device())
rusty1s's avatar
rusty1s committed
440
            else:
rusty1s's avatar
rusty1s committed
441
                value = torch.ones(self.nnz(), device=self.device())
rusty1s's avatar
rusty1s committed
442

rusty1s's avatar
rusty1s committed
443
        return torch.sparse_coo_tensor(index, value, self.sizes())
rusty1s's avatar
rusty1s committed
444

rusty1s's avatar
rusty1s committed
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459

# Python Bindings #############################################################

Dtype = Optional[torch.dtype]
Device = Optional[Union[torch.device, str]]


def share_memory_(self: SparseTensor) -> SparseTensor:
    self.storage.share_memory_()


def is_shared(self: SparseTensor) -> bool:
    return self.storage.is_shared()


rusty1s's avatar
typing  
rusty1s committed
460
461
462
def to(self, *args: Optional[List[Any]],
       **kwargs: Optional[Dict[str, Any]]) -> SparseTensor:

rusty1s's avatar
rusty1s committed
463
    device, dtype, non_blocking = torch._C._nn._parse_to(*args, **kwargs)[:3]
rusty1s's avatar
rusty1s committed
464
465
466
467
468
469
470
471
472

    if dtype is not None:
        self = self.type_as(torch.tensor(0., dtype=dtype))
    if device is not None:
        self = self.device_as(torch.tensor(0., device=device), non_blocking)

    return self


rusty1s's avatar
typing  
rusty1s committed
473
def __getitem__(self: SparseTensor, index: Any) -> SparseTensor:
rusty1s's avatar
repr  
rusty1s committed
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
511
512
513
514
    index = list(index) if isinstance(index, tuple) else [index]
    # More than one `Ellipsis` is not allowed...
    if len([i for i in index if not torch.is_tensor(i) and i == ...]) > 1:
        raise SyntaxError

    dim = 0
    out = self
    while len(index) > 0:
        item = index.pop(0)
        if isinstance(item, int):
            out = out.select(dim, item)
            dim += 1
        elif isinstance(item, slice):
            if item.step is not None:
                raise ValueError('Step parameter not yet supported.')

            start = 0 if item.start is None else item.start
            start = self.size(dim) + start if start < 0 else start

            stop = self.size(dim) if item.stop is None else item.stop
            stop = self.size(dim) + stop if stop < 0 else stop

            out = out.narrow(dim, start, max(stop - start, 0))
            dim += 1
        elif torch.is_tensor(item):
            if item.dtype == torch.bool:
                out = out.masked_select(dim, item)
                dim += 1
            elif item.dtype == torch.long:
                out = out.index_select(dim, item)
                dim += 1
        elif item == Ellipsis:
            if self.dim() - len(index) < dim:
                raise SyntaxError
            dim = self.dim() - len(index)
        else:
            raise SyntaxError

    return out


rusty1s's avatar
typing  
rusty1s committed
515
def __repr__(self: SparseTensor) -> str:
rusty1s's avatar
repr  
rusty1s committed
516
517
518
    i = ' ' * 6
    row, col, value = self.coo()
    infos = []
rusty1s's avatar
rusty1s committed
519
520
    infos += [f'row={indent(row.__repr__(), i)[len(i):]}']
    infos += [f'col={indent(col.__repr__(), i)[len(i):]}']
rusty1s's avatar
repr  
rusty1s committed
521
522

    if value is not None:
rusty1s's avatar
rusty1s committed
523
        infos += [f'val={indent(value.__repr__(), i)[len(i):]}']
rusty1s's avatar
repr  
rusty1s committed
524
525

    infos += [
rusty1s's avatar
rusty1s committed
526
527
        f'size={tuple(self.sizes())}, nnz={self.nnz()}, '
        f'density={100 * self.density():.02f}%'
rusty1s's avatar
repr  
rusty1s committed
528
    ]
rusty1s's avatar
rusty1s committed
529

rusty1s's avatar
repr  
rusty1s committed
530
531
532
    infos = ',\n'.join(infos)

    i = ' ' * (len(self.__class__.__name__) + 1)
rusty1s's avatar
rusty1s committed
533
    return f'{self.__class__.__name__}({indent(infos, i)[len(i):]})'
rusty1s's avatar
repr  
rusty1s committed
534
535


rusty1s's avatar
rusty1s committed
536
537
538
SparseTensor.share_memory_ = share_memory_
SparseTensor.is_shared = is_shared
SparseTensor.to = to
rusty1s's avatar
repr  
rusty1s committed
539
540
SparseTensor.__getitem__ = __getitem__
SparseTensor.__repr__ = __repr__
rusty1s's avatar
rusty1s committed
541
542
543

# Scipy Conversions ###########################################################

rusty1s's avatar
rusty1s committed
544
545
ScipySparseMatrix = Union[scipy.sparse.coo_matrix, scipy.sparse.csr_matrix,
                          scipy.sparse.csc_matrix]
rusty1s's avatar
rusty1s committed
546
547
548


@torch.jit.ignore
rusty1s's avatar
rusty1s committed
549
def from_scipy(mat: ScipySparseMatrix, has_value: bool = True) -> SparseTensor:
rusty1s's avatar
rusty1s committed
550
551
552
553
554
555
556
557
558
    colptr = None
    if isinstance(mat, scipy.sparse.csc_matrix):
        colptr = torch.from_numpy(mat.indptr).to(torch.long)

    mat = mat.tocsr()
    rowptr = torch.from_numpy(mat.indptr).to(torch.long)
    mat = mat.tocoo()
    row = torch.from_numpy(mat.row).to(torch.long)
    col = torch.from_numpy(mat.col).to(torch.long)
rusty1s's avatar
rusty1s committed
559
560
561
    value = None
    if has_value:
        value = torch.from_numpy(mat.data)
rusty1s's avatar
rusty1s committed
562
563
    sparse_sizes = mat.shape[:2]

rusty1s's avatar
rusty1s committed
564
565
566
567
    storage = SparseStorage(row=row, rowptr=rowptr, col=col, value=value,
                            sparse_sizes=sparse_sizes, rowcount=None,
                            colptr=colptr, colcount=None, csr2csc=None,
                            csc2csr=None, is_sorted=True)
rusty1s's avatar
rusty1s committed
568
569
570
571
572

    return SparseTensor.from_storage(storage)


@torch.jit.ignore
rusty1s's avatar
rusty1s committed
573
def to_scipy(self: SparseTensor, layout: Optional[str] = None,
rusty1s's avatar
rusty1s committed
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
             dtype: Optional[torch.dtype] = None) -> ScipySparseMatrix:
    assert self.dim() == 2
    layout = get_layout(layout)

    if not self.has_value():
        ones = torch.ones(self.nnz(), dtype=dtype).numpy()

    if layout == 'coo':
        row, col, value = self.coo()
        row = row.detach().cpu().numpy()
        col = col.detach().cpu().numpy()
        value = value.detach().cpu().numpy() if self.has_value() else ones
        return scipy.sparse.coo_matrix((value, (row, col)), self.sizes())
    elif layout == 'csr':
        rowptr, col, value = self.csr()
        rowptr = rowptr.detach().cpu().numpy()
        col = col.detach().cpu().numpy()
        value = value.detach().cpu().numpy() if self.has_value() else ones
        return scipy.sparse.csr_matrix((value, col, rowptr), self.sizes())
    elif layout == 'csc':
        colptr, row, value = self.csc()
        colptr = colptr.detach().cpu().numpy()
        row = row.detach().cpu().numpy()
        value = value.detach().cpu().numpy() if self.has_value() else ones
        return scipy.sparse.csc_matrix((value, row, colptr), self.sizes())


SparseTensor.from_scipy = from_scipy
SparseTensor.to_scipy = to_scipy

# Hacky fixes #################################################################

rusty1s's avatar
rusty1s committed
606
# Fix standard operators of `torch.Tensor` for PyTorch<=1.3.
rusty1s's avatar
rusty1s committed
607
# https://github.com/pytorch/pytorch/pull/31769
rusty1s's avatar
rusty1s committed
608
609
TORCH_MAJOR = int(torch.__version__.split('.')[0])
TORCH_MINOR = int(torch.__version__.split('.')[1])
rusty1s's avatar
rusty1s committed
610
if (TORCH_MAJOR < 1) or (TORCH_MAJOR == 1 and TORCH_MINOR <= 3):
rusty1s's avatar
rusty1s committed
611
612

    def add(self, other):
rusty1s's avatar
rusty1s committed
613
614
615
        if torch.is_tensor(other) or is_scalar(other):
            return self.add(other)
        return NotImplemented
rusty1s's avatar
rusty1s committed
616
617

    def mul(self, other):
rusty1s's avatar
rusty1s committed
618
619
620
        if torch.is_tensor(other) or is_scalar(other):
            return self.mul(other)
        return NotImplemented
rusty1s's avatar
rusty1s committed
621
622

    torch.Tensor.__add__ = add
rusty1s's avatar
rusty1s committed
623
    torch.Tensor.__mul__ = mul