narrow.py 2.63 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
import torch
rusty1s's avatar
rusty1s committed
2
3
from torch_sparse.storage import SparseStorage
from torch_sparse.tensor import SparseTensor
rusty1s's avatar
rusty1s committed
4
5


rusty1s's avatar
rusty1s committed
6
7
@torch.jit.script
def narrow(src: SparseTensor, dim: int, start: int, length: int):
rusty1s's avatar
rusty1s committed
8
9
    dim = src.dim() + dim if dim < 0 else dim
    start = src.size(dim) + start if start < 0 else start
rusty1s's avatar
rusty1s committed
10

rusty1s's avatar
rusty1s committed
11
    if dim == 0:
rusty1s's avatar
rusty1s committed
12
        rowptr, col, value = src.csr()
rusty1s's avatar
rusty1s committed
13

rusty1s's avatar
rusty1s committed
14
15
16
17
        rowptr = rowptr.narrow(0, start=start, length=length + 1)
        row_start = rowptr[0]
        rowptr = rowptr - row_start
        row_length = rowptr[-1]
rusty1s's avatar
rusty1s committed
18

rusty1s's avatar
rusty1s committed
19
20
21
        row = src.storage._row
        if row is not None:
            row = row.narrow(0, row_start, row_length) - start
rusty1s's avatar
rusty1s committed
22

rusty1s's avatar
rusty1s committed
23
        col = col.narrow(0, row_start, row_length)
rusty1s's avatar
rusty1s committed
24

rusty1s's avatar
rusty1s committed
25
        if value is not None:
rusty1s's avatar
rusty1s committed
26
            value = value.narrow(0, row_start, row_length)
rusty1s's avatar
rusty1s committed
27

rusty1s's avatar
rusty1s committed
28
        sparse_sizes = torch.Size([length, src.sparse_size(1)])
rusty1s's avatar
rusty1s committed
29

rusty1s's avatar
rusty1s committed
30
31
32
33
34
35
36
37
38
        rowcount = src.storage._rowcount
        if rowcount is not None:
            rowcount = rowcount.narrow(0, start=start, length=length)

        storage = SparseStorage(row=row, rowptr=rowptr, col=col, value=value,
                                sparse_sizes=sparse_sizes, rowcount=rowcount,
                                colptr=None, colcount=None, csr2csc=None,
                                csc2csr=None, is_sorted=True)
        return src.from_storage(storage)
rusty1s's avatar
rusty1s committed
39

rusty1s's avatar
rusty1s committed
40
    elif dim == 1:
rusty1s's avatar
rusty1s committed
41
        # This is faster than accessing `csc()` contrary to the `dim=0` case.
rusty1s's avatar
rusty1s committed
42
        row, col, value = src.coo()
rusty1s's avatar
rusty1s committed
43
        mask = (col >= start) & (col < start + length)
rusty1s's avatar
rusty1s committed
44

rusty1s's avatar
rusty1s committed
45
46
        row = row[mask]
        col = col[mask] - start
rusty1s's avatar
rusty1s committed
47

rusty1s's avatar
rusty1s committed
48
49
50
51
        if value is not None:
            value = value[mask]

        sparse_sizes = torch.Size([src.sparse_size(0), length])
rusty1s's avatar
rusty1s committed
52
53

        colptr = src.storage._colptr
rusty1s's avatar
rusty1s committed
54
55
56
57
        if colptr is not None:
            colptr = colptr.narrow(0, start=start, length=length + 1)
            colptr = colptr - colptr[0]

rusty1s's avatar
rusty1s committed
58
59
60
        colcount = src.storage._colcount
        if colcount is not None:
            colcount = colcount.narrow(0, start=start, length=length)
rusty1s's avatar
rusty1s committed
61

rusty1s's avatar
rusty1s committed
62
63
64
65
66
        storage = SparseStorage(row=row, rowptr=None, col=col, value=value,
                                sparse_sizes=sparse_sizes, rowcount=None,
                                colptr=colptr, colcount=colcount, csr2csc=None,
                                csc2csr=None, is_sorted=True)
        return src.from_storage(storage)
rusty1s's avatar
rusty1s committed
67
68

    else:
rusty1s's avatar
rusty1s committed
69
70
71
72
73
74
75
        value = src.storage.value()
        if value is not None:
            return src.set_value(value.narrow(dim - 1, start, length),
                                 layout='coo')
        else:
            raise ValueError

rusty1s's avatar
rusty1s committed
76

rusty1s's avatar
rusty1s committed
77
78
SparseTensor.narrow = lambda self, dim, start, length: narrow(
    self, dim, start, length)