tensor.py 17.9 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
2
3
4
5
from textwrap import indent

import torch
import scipy.sparse

rusty1s's avatar
rusty1s committed
6
from torch_sparse.storage import SparseStorage, get_layout
rusty1s's avatar
rusty1s committed
7

rusty1s's avatar
rusty1s committed
8
from torch_sparse.transpose import t
rusty1s's avatar
rusty1s committed
9
from torch_sparse.narrow import narrow
rusty1s's avatar
rusty1s committed
10
11
12
from torch_sparse.select import select
from torch_sparse.index_select import index_select, index_select_nnz
from torch_sparse.masked_select import masked_select, masked_select_nnz
rusty1s's avatar
rusty1s committed
13
14
15
16


class SparseTensor(object):
    def __init__(self, index, value=None, sparse_size=None, is_sorted=False):
rusty1s's avatar
rusty1s committed
17
18
        self.storage = SparseStorage(
            index, value, sparse_size, is_sorted=is_sorted)
rusty1s's avatar
rusty1s committed
19
20
21
22

    @classmethod
    def from_storage(self, storage):
        self = SparseTensor.__new__(SparseTensor)
rusty1s's avatar
rusty1s committed
23
        self.storage = storage
rusty1s's avatar
rusty1s committed
24
25
26
27
28
29
30
31
32
33
34
        return self

    @classmethod
    def from_dense(self, mat):
        if mat.dim() > 2:
            index = mat.abs().sum([i for i in range(2, mat.dim())]).nonzero()
        else:
            index = mat.nonzero()

        index = index.t().contiguous()
        value = mat[index[0], index[1]]
rusty1s's avatar
rusty1s committed
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
        return SparseTensor(index, value, mat.size()[:2], is_sorted=True)

    @classmethod
    def from_torch_sparse_coo_tensor(self, mat, is_sorted=False):
        return SparseTensor(
            mat._indices(), mat._values(), mat.size()[:2], is_sorted=is_sorted)

    @classmethod
    def from_scipy(self, mat):
        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)
        index = torch.stack([row, col], dim=0)
        value = torch.from_numpy(mat.data)
        size = mat.shape

        storage = SparseStorage(
            index, value, size, rowptr=rowptr, colptr=colptr, is_sorted=True)

        return SparseTensor.from_storage(storage)
rusty1s's avatar
rusty1s committed
61
62

    def __copy__(self):
rusty1s's avatar
rusty1s committed
63
        return self.from_storage(self.storage)
rusty1s's avatar
rusty1s committed
64
65

    def clone(self):
rusty1s's avatar
rusty1s committed
66
        return self.from_storage(self.storage.clone())
rusty1s's avatar
rusty1s committed
67
68
69
70
71
72
73
74
75

    def __deepcopy__(self, memo):
        new_sparse_tensor = self.clone()
        memo[id(self)] = new_sparse_tensor
        return new_sparse_tensor

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

    def coo(self):
rusty1s's avatar
rusty1s committed
76
        return self.storage.index, self.storage.value
rusty1s's avatar
rusty1s committed
77
78

    def csr(self):
rusty1s's avatar
rusty1s committed
79
        return self.storage.rowptr, self.storage.col, self.storage.value
rusty1s's avatar
rusty1s committed
80
81

    def csc(self):
rusty1s's avatar
rusty1s committed
82
83
84
        perm = self.storage.csr2csc
        return (self.storage.colptr, self.storage.row[perm],
                self.storage.value[perm] if self.has_value() else None)
rusty1s's avatar
rusty1s committed
85
86
87
88

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

    def has_value(self):
rusty1s's avatar
rusty1s committed
89
        return self.storage.has_value()
rusty1s's avatar
rusty1s committed
90
91

    def set_value_(self, value, layout=None):
rusty1s's avatar
rusty1s committed
92
        self.storage.set_value_(value, layout)
rusty1s's avatar
rusty1s committed
93
94
95
        return self

    def set_value(self, value, layout=None):
rusty1s's avatar
rusty1s committed
96
        return self.from_storage(self.storage.set_value(value, layout))
rusty1s's avatar
rusty1s committed
97
98

    def sparse_size(self, dim=None):
rusty1s's avatar
rusty1s committed
99
        return self.storage.sparse_size(dim)
rusty1s's avatar
rusty1s committed
100
101

    def sparse_resize_(self, *sizes):
rusty1s's avatar
rusty1s committed
102
        self.storage.sparse_resize_(*sizes)
rusty1s's avatar
rusty1s committed
103
104
105
        return self

    def is_coalesced(self):
rusty1s's avatar
rusty1s committed
106
        return self.storage.is_coalesced()
rusty1s's avatar
rusty1s committed
107
108

    def coalesce(self):
rusty1s's avatar
rusty1s committed
109
        return self.from_storage(self.storage.coalesce())
rusty1s's avatar
rusty1s committed
110
111

    def cached_keys(self):
rusty1s's avatar
rusty1s committed
112
        return self.storage.cached_keys()
rusty1s's avatar
rusty1s committed
113
114

    def fill_cache_(self, *args):
rusty1s's avatar
rusty1s committed
115
        self.storage.fill_cache_(*args)
rusty1s's avatar
rusty1s committed
116
117
118
        return self

    def clear_cache_(self, *args):
rusty1s's avatar
rusty1s committed
119
        self.storage.clear_cache_(*args)
rusty1s's avatar
rusty1s committed
120
121
122
123
124
125
        return self

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

    def size(self, dim=None):
        size = self.sparse_size()
rusty1s's avatar
rusty1s committed
126
        size += self.storage.value.size()[1:] if self.has_value() else ()
rusty1s's avatar
rusty1s committed
127
128
129
130
131
132
133
134
135
136
        return size if dim is None else size[dim]

    def dim(self):
        return len(self.size())

    @property
    def shape(self):
        return self.size()

    def nnz(self):
rusty1s's avatar
rusty1s committed
137
        return self.storage.index.size(1)
rusty1s's avatar
rusty1s committed
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162

    def density(self):
        return self.nnz() / (self.sparse_size(0) * self.sparse_size(1))

    def sparsity(self):
        return 1 - self.density()

    def avg_row_length(self):
        return self.nnz() / self.sparse_size(0)

    def avg_col_length(self):
        return self.nnz() / self.sparse_size(1)

    def numel(self):
        return self.value.numel() if self.has_value() else self.nnz()

    def is_quadratic(self):
        return self.sparse_size(0) == self.sparse_size(1)

    def is_symmetric(self):
        if not self.is_quadratic:
            return False

        rowptr, col, val1 = self.csr()
        colptr, row, val2 = self.csc()
rusty1s's avatar
rusty1s committed
163
164
165
        index_sym = (rowptr == colptr).all() and (col == row).all()
        value_sym = (val1 == val2).all().item() if self.has_value() else True
        return index_sym.item() and value_sym
rusty1s's avatar
rusty1s committed
166
167

    def detach_(self):
rusty1s's avatar
rusty1s committed
168
        self.storage.apply_(lambda x: x.detach_())
rusty1s's avatar
rusty1s committed
169
170
171
        return self

    def detach(self):
rusty1s's avatar
rusty1s committed
172
        return self.from_storage(self.storage.apply(lambda x: x.detach()))
rusty1s's avatar
rusty1s committed
173
174

    def pin_memory(self):
rusty1s's avatar
rusty1s committed
175
        return self.from_storage(self.storage.apply(lambda x: x.pin_memory()))
rusty1s's avatar
rusty1s committed
176
177

    def is_pinned(self):
rusty1s's avatar
rusty1s committed
178
        return all(self.storage.map(lambda x: x.is_pinned()))
rusty1s's avatar
rusty1s committed
179
180

    def share_memory_(self):
rusty1s's avatar
rusty1s committed
181
        self.storage.apply_(lambda x: x.share_memory_())
rusty1s's avatar
rusty1s committed
182
183
184
        return self

    def is_shared(self):
rusty1s's avatar
rusty1s committed
185
        return all(self.storage.map(lambda x: x.is_shared()))
rusty1s's avatar
rusty1s committed
186
187
188

    @property
    def device(self):
rusty1s's avatar
rusty1s committed
189
        return self.storage.index.device
rusty1s's avatar
rusty1s committed
190
191

    def cpu(self):
rusty1s's avatar
rusty1s committed
192
        return self.from_storage(self.storage.apply(lambda x: x.cpu()))
rusty1s's avatar
rusty1s committed
193
194

    def cuda(self, device=None, non_blocking=False, **kwargs):
rusty1s's avatar
rusty1s committed
195
196
        storage = self.storage.apply(lambda x: x.cuda(device, non_blocking, **
                                                      kwargs))
rusty1s's avatar
rusty1s committed
197
        return self.from_storage(storage)
rusty1s's avatar
rusty1s committed
198
199
200

    @property
    def is_cuda(self):
rusty1s's avatar
rusty1s committed
201
        return self.storage.index.is_cuda
rusty1s's avatar
rusty1s committed
202
203
204

    @property
    def dtype(self):
rusty1s's avatar
rusty1s committed
205
        return self.storage.value.dtype if self.has_value() else None
rusty1s's avatar
rusty1s committed
206
207

    def is_floating_point(self):
rusty1s's avatar
rusty1s committed
208
        value = self.storage.value
rusty1s's avatar
rusty1s committed
209
210
211
212
213
214
215
216
217
        return self.has_value() and torch.is_floating_point(value)

    def type(self, dtype=None, non_blocking=False, **kwargs):
        if dtype is None:
            return self.dtype

        if dtype == self.dtype:
            return self

rusty1s's avatar
rusty1s committed
218
219
        storage = self.storage.apply_value(lambda x: x.type(
            dtype, non_blocking, **kwargs))
rusty1s's avatar
rusty1s committed
220
221

        return self.from_storage(storage)
rusty1s's avatar
rusty1s committed
222
223
224
225
226
227
228

    def to(self, *args, **kwargs):
        storage = None

        if 'device' in kwargs:
            device = kwargs['device']
            del kwargs['device']
rusty1s's avatar
rusty1s committed
229
            storage = self.storage.apply(lambda x: x.to(
rusty1s's avatar
rusty1s committed
230
231
232
233
                device, non_blocking=getattr(kwargs, 'non_blocking', False)))

        for arg in args[:]:
            if isinstance(arg, str) or isinstance(arg, torch.device):
rusty1s's avatar
rusty1s committed
234
                storage = self.storage.apply(lambda x: x.to(
rusty1s's avatar
rusty1s committed
235
236
237
238
                    arg, non_blocking=getattr(kwargs, 'non_blocking', False)))
                args.remove(arg)

        if storage is not None:
rusty1s's avatar
rusty1s committed
239
            self = self.from_storage(storage)
rusty1s's avatar
rusty1s committed
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287

        if len(args) > 0 or len(kwargs) > 0:
            self = self.type(*args, **kwargs)

        return self

    def bfloat16(self):
        return self.type(torch.bfloat16)

    def bool(self):
        return self.type(torch.bool)

    def byte(self):
        return self.type(torch.byte)

    def char(self):
        return self.type(torch.char)

    def half(self):
        return self.type(torch.half)

    def float(self):
        return self.type(torch.float)

    def double(self):
        return self.type(torch.double)

    def short(self):
        return self.type(torch.short)

    def int(self):
        return self.type(torch.int)

    def long(self):
        return self.type(torch.long)

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

    def to_dense(self, dtype=None):
        dtype = dtype or self.dtype
        (row, col), value = self.coo()
        mat = torch.zeros(self.size(), dtype=dtype, device=self.device)
        mat[row, col] = value if self.has_value() else 1
        return mat

    def to_torch_sparse_coo_tensor(self, dtype=None, requires_grad=False):
        index, value = self.coo()
        return torch.sparse_coo_tensor(
rusty1s's avatar
rusty1s committed
288
289
290
291
292
293
            index,
            value if self.has_value() else torch.ones(
                self.nnz(), dtype=dtype, device=self.device),
            self.size(),
            device=self.device,
            requires_grad=requires_grad)
rusty1s's avatar
rusty1s committed
294
295

    def to_scipy(self, dtype=None, layout=None):
rusty1s's avatar
rusty1s committed
296
        assert self.dim() == 2
rusty1s's avatar
rusty1s committed
297
        layout = get_layout(layout)
rusty1s's avatar
rusty1s committed
298

rusty1s's avatar
rusty1s committed
299
300
        if not self.has_value():
            ones = torch.ones(self.nnz(), dtype=dtype).numpy()
rusty1s's avatar
rusty1s committed
301
302

        if layout == 'coo':
rusty1s's avatar
rusty1s committed
303
304
305
306
            (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
rusty1s's avatar
rusty1s committed
307
308
            return scipy.sparse.coo_matrix((value, (row, col)), self.size())
        elif layout == 'csr':
rusty1s's avatar
rusty1s committed
309
310
311
312
            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
rusty1s's avatar
rusty1s committed
313
314
            return scipy.sparse.csr_matrix((value, col, rowptr), self.size())
        elif layout == 'csc':
rusty1s's avatar
rusty1s committed
315
316
317
318
            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
rusty1s's avatar
rusty1s committed
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
            return scipy.sparse.csc_matrix((value, row, colptr), self.size())

    # String Reputation #######################################################

    def __repr__(self):
        i = ' ' * 6
        index, value = self.coo()
        infos = [f'index={indent(index.__repr__(), i)[len(i):]}']

        if self.has_value():
            infos += [f'value={indent(value.__repr__(), i)[len(i):]}']

        infos += [
            f'size={tuple(self.size())}, '
            f'nnz={self.nnz()}, '
            f'density={100 * self.density():.02f}%'
        ]
        infos = ',\n'.join(infos)

        i = ' ' * (len(self.__class__.__name__) + 1)
        return f'{self.__class__.__name__}({indent(infos, i)[len(i):]})'


# Bindings ####################################################################
rusty1s's avatar
rusty1s committed
343

rusty1s's avatar
rusty1s committed
344
SparseTensor.t = t
rusty1s's avatar
rusty1s committed
345
SparseTensor.narrow = narrow
rusty1s's avatar
rusty1s committed
346
347
348
349
350
SparseTensor.select = select
SparseTensor.index_select = index_select
SparseTensor.index_select_nnz = index_select_nnz
SparseTensor.masked_select = masked_select
SparseTensor.masked_select_nnz = masked_select_nnz
rusty1s's avatar
rusty1s committed
351
352
353
354
355
356

#     def __getitem__(self, idx):
#         # Convert int and slice to index tensor
#         # Filter list into edge and sparse slice
#         raise NotImplementedError

rusty1s's avatar
typo  
rusty1s committed
357
358
359
# def remove_diag(self):
#     raise NotImplementedError

rusty1s's avatar
rusty1s committed
360
361
362
#     def set_diag(self, value):
#         raise NotImplementedError

rusty1s's avatar
rusty1s committed
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
#     def __reduce(self, dim, reduce, only_nnz):
#         raise NotImplementedError

#     def sum(self, dim):
#         return self.__reduce(dim, reduce='add', only_nnz=True)

#     def prod(self, dim):
#         return self.__reduce(dim, reduce='mul', only_nnz=True)

#     def min(self, dim, only_nnz=False):
#         return self.__reduce(dim, reduce='min', only_nnz=only_nnz)

#     def max(self, dim, only_nnz=False):
#         return self.__reduce(dim, reduce='min', only_nnz=only_nnz)

#     def mean(self, dim, only_nnz=False):
#         return self.__reduce(dim, reduce='mean', only_nnz=only_nnz)

#     def matmul(self, mat, reduce='add'):
#         assert self.numel() == self.nnz()  # Disallow multi-dimensional value
#         if torch.is_tensor(mat):
#             raise NotImplementedError
#         elif isinstance(mat, self.__class__):
#             assert reduce == 'add'
rusty1s's avatar
rusty1s committed
387
#           assert mat.numel() == mat.nnz()  # Disallow multi-dimensional value
rusty1s's avatar
rusty1s committed
388
389
390
391
#             raise NotImplementedError
#         raise ValueError('Argument needs to be of type `torch.tensor` or '
#                          'type `torch_sparse.SparseTensor`.')

rusty1s's avatar
typo  
rusty1s committed
392
393
# def add_nnz(self):

rusty1s's avatar
rusty1s committed
394
395
396
397
398
399
400
401
402
403
#     def add(self, other, layout=None):
#         if __is_scalar__(other):
#             if self.has_value:
#                 return self.set_value(self._value + other, 'coo')
#             else:
#                 return self.set_value(torch.full((self.nnz(), ), other + 1),
#                                       'coo')
#         elif torch.is_tensor(other):
#             if layout is None:
#                 layout = 'coo'
rusty1s's avatar
rusty1s committed
404
405
#               warnings.warn('`layout` argument unset, using default layout '
#                             '"coo". This may lead to unexpected behaviour.')
rusty1s's avatar
rusty1s committed
406
407
#             assert layout in ['coo', 'csr', 'csc']
#             if layout == 'csc':
rusty1s's avatar
rusty1s committed
408
#                 other = other[self._arg_csc2csr]
rusty1s's avatar
rusty1s committed
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
#             if self.has_value:
#                 return self.set_value(self._value + other, 'coo')
#             else:
#                 return self.set_value(other + 1, 'coo')
#         elif isinstance(other, self.__class__):
#             raise NotImplementedError
#         raise ValueError('Argument needs to be of type `int`, `float`, '
#                          '`torch.tensor` or `torch_sparse.SparseTensor`.')

#     def add_(self, other, layout=None):
#         if isinstance(other, int) or isinstance(other, float):
#             raise NotImplementedError
#         elif torch.is_tensor(other):
#             raise NotImplementedError
#         raise ValueError('Argument needs to be a scalar or of type '
#                          '`torch.tensor`.')

#     def __add__(self, other):
#         return self.add(other)

#     def __radd__(self, other):
#         return self.add(other)

#     def sub(self, layout=None):
#         raise NotImplementedError

#     def sub_(self, layout=None):
#         raise NotImplementedError

#     def mul(self, layout=None):
#         raise NotImplementedError

#     def mul_(self, layout=None):
#         raise NotImplementedError

#     def div(self, layout=None):
#         raise NotImplementedError

#     def div_(self, layout=None):
#         raise NotImplementedError

rusty1s's avatar
rusty1s committed
450
451
452
if __name__ == '__main__':
    from torch_geometric.datasets import Reddit, Planetoid  # noqa
    import time  # noqa
rusty1s's avatar
rusty1s committed
453

rusty1s's avatar
rusty1s committed
454
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
rusty1s's avatar
rusty1s committed
455

rusty1s's avatar
rusty1s committed
456
457
458
    # dataset = Reddit('/tmp/Reddit')
    dataset = Planetoid('/tmp/Cora', 'Cora')
    data = dataset[0].to(device)
rusty1s's avatar
rusty1s committed
459

rusty1s's avatar
rusty1s committed
460
    value = torch.randn((data.num_edges, ), device=device)
rusty1s's avatar
rusty1s committed
461

rusty1s's avatar
rusty1s committed
462
    mat1 = SparseTensor(data.edge_index, value)
rusty1s's avatar
rusty1s committed
463

rusty1s's avatar
rusty1s committed
464
465
466
467
468
469
470
471
    mat1 = SparseTensor.from_dense(mat1.to_dense())

    print(mat1)
    mat = SparseTensor.from_torch_sparse_coo_tensor(
        mat1.to_torch_sparse_coo_tensor())

    mat = SparseTensor.from_scipy(mat.to_scipy(layout='csc'))
    print(mat)
rusty1s's avatar
rusty1s committed
472

rusty1s's avatar
rusty1s committed
473
474
    # index = torch.tensor([0, 2])
    # mat2 = mat1.index_select(2, index)
rusty1s's avatar
rusty1s committed
475

rusty1s's avatar
rusty1s committed
476
477
478
    # index = torch.randperm(data.num_nodes)[:data.num_nodes - 500]
    # mask = torch.zeros(data.num_nodes, dtype=torch.bool)
    # mask[index] = True
rusty1s's avatar
rusty1s committed
479

rusty1s's avatar
rusty1s committed
480
481
482
483
484
485
486
487
488
    # t = time.perf_counter()
    # for _ in range(1000):
    #     mat2 = mat1.index_select(0, index)
    # print(time.perf_counter() - t)

    # t = time.perf_counter()
    # for _ in range(1000):
    #     mat2 = mat1.masked_select(0, mask)
    # print(time.perf_counter() - t)
rusty1s's avatar
rusty1s committed
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509

    # mat2 = mat1.narrow(1, start=0, length=3)
    # print(mat2)

    # index = torch.randperm(data.num_nodes)
    # t = time.perf_counter()
    # for _ in range(1000):
    #     mat2 = mat1.index_select(0, index)
    # print(time.perf_counter() - t)

    # t = time.perf_counter()
    # for _ in range(1000):
    #     mat2 = mat1.index_select(1, index)
    # print(time.perf_counter() - t)
    # raise NotImplementedError

    # t = time.perf_counter()
    # for _ in range(1000):
    #     mat2 = mat1.t().index_select(0, index).t()
    # print(time.perf_counter() - t)

rusty1s's avatar
rusty1s committed
510
    # print(mat1)
rusty1s's avatar
rusty1s committed
511
512
513
514
515
516
517
518
519
520
521
522
523
    # mat1.index_select((0, 1), torch.tensor([0, 1, 2, 3, 4]))

    # print(mat3)
    # print(mat3.storage.rowcount)

    # print(mat1)

    # (row, col), value = mat1.coo()
    # mask = row < 3
    # t = time.perf_counter()
    # for _ in range(10000):
    #     mat2 = mat1.narrow(1, start=10, length=2690)
    # print(time.perf_counter() - t)
rusty1s's avatar
rusty1s committed
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540

    # # print(mat1.to_dense().size())
    # print(mat1.to_torch_sparse_coo_tensor().to_dense().size())
    # print(mat1.to_scipy(layout='coo').todense().shape)
    # print(mat1.to_scipy(layout='csr').todense().shape)
    # print(mat1.to_scipy(layout='csc').todense().shape)

    # print(mat1.is_quadratic())
    # print(mat1.is_symmetric())

    # print(mat1.cached_keys())
    # mat1 = mat1.t()
    # print(mat1.cached_keys())
    # mat1 = mat1.t()
    # print(mat1.cached_keys())
    # print('-------- NARROW ----------')

rusty1s's avatar
rusty1s committed
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
    # t = time.perf_counter()
    # for _ in range(100):
    #     out = mat1.narrow(dim=0, start=10, length=10)
    #     # out.storage.colptr
    # print(time.perf_counter() - t)
    # print(out)
    # print(out.cached_keys())

    # t = time.perf_counter()
    # for _ in range(100):
    #     out = mat1.narrow(dim=1, start=10, length=2000)
    #     # out.storage.colptr
    # print(time.perf_counter() - t)
    # print(out)
    # print(out.cached_keys())
rusty1s's avatar
rusty1s committed
556
557

    # mat1 = mat1.narrow(0, start=10, length=10)
rusty1s's avatar
rusty1s committed
558
    # mat1.storage._value = torch.randn(mat1.nnz(), 20)
rusty1s's avatar
rusty1s committed
559
560
561
    # print(mat1.coo()[1].size())
    # mat1 = mat1.narrow(2, start=10, length=10)
    # print(mat1.coo()[1].size())
rusty1s's avatar
rusty1s committed
562
563
#     mat1 = mat1.t()

rusty1s's avatar
rusty1s committed
564
#   mat2 = torch.sparse_coo_tensor(data.edge_index, torch.ones(data.num_edges),
rusty1s's avatar
rusty1s committed
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
#                                    device=device)
#     mat2 = mat2.coalesce()
#     mat2 = mat2.t().coalesce()

#     index1, value1 = mat1.coo()
#     index2, value2 = mat2._indices(), mat2._values()
#     assert torch.allclose(index1, index2)

#     out1 = mat1.to_dense()
#     out2 = mat2.to_dense()
#     assert torch.allclose(out1, out2)

#     out = 2 + mat1
#     print(out)

#     # mat1[1]
#     # mat1[1, 1]
#     # mat1[..., -1]
#     # mat1[:, -1]
#     # mat1[1:4, 1:4]
#     # mat1[torch.tensor([0, 1, 2])]