Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
test_jit.py 2.06 KB
Newer Older
rusty1s's avatar
beginning of torch script support  
rusty1s committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
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
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80 
import torch

from torch_sparse.tensor import SparseTensor
from torch_sparse.storage import SparseStorage

from typing import Dict, Any

# class MyTensor(dict):
#     def __init__(self, rowptr, col):
#         self['rowptr'] = rowptr
#         self['col'] = col

# def rowptr(self: Dict[str, torch.Tensor]):
#     return self['rowptr']


@torch.jit.script
class Foo:
    rowptr: torch.Tensor
    col: torch.Tensor

    def __init__(self, rowptr: torch.Tensor, col: torch.Tensor):
        self.rowptr = rowptr
        self.col = col


class MyCell(torch.nn.Module):
    def __init__(self):
        super(MyCell, self).__init__()
        self.linear = torch.nn.Linear(2, 4)

    # def forward(self, x: torch.Tensor, ptr: torch.Tensor) -> torch.Tensor:
    def forward(self, x: torch.Tensor, adj: SparseStorage) -> torch.Tensor:
        out, _ = torch.ops.torch_sparse_cpu.spmm(adj.rowptr(), adj.col(), None,
                                                 x, 'sum')
        return out


#         ind = torch.ops.torch_sparse_cpu.ptr2ind(ptr, ptr[-1].item())
#         # ind = ptr2ind(ptr, E)
#         x_j = x[ind]
#         out = self.linear(x_j)
#         return out


def test_jit():
    my_cell = MyCell()

    # x = torch.rand(3, 2)
    # ptr = torch.tensor([0, 2, 4, 6])
    # out = my_cell(x, ptr)
    # print()
    # print(out)

    # traced_cell = torch.jit.trace(my_cell, (x, ptr))
    # print(traced_cell)
    # out = traced_cell(x, ptr)
    # print(out)

    x = torch.randn(3, 2)

    # adj = torch.randn(3, 3)
    # adj = SparseTensor.from_dense(adj)
    # adj = Foo(adj.storage.rowptr, adj.storage.col)
    # adj = adj.storage

    rowptr = torch.tensor([0, 3, 6, 9])
    col = torch.tensor([0, 1, 2, 0, 1, 2, 0, 1, 2])

    adj = SparseStorage(rowptr=rowptr, col=col)

    # adj = {'rowptr': mat.storage.rowptr, 'col': mat.storage.col}
    # foo = Foo(mat.storage.rowptr, mat.storage.col)
    # adj = MyTensor(mat.storage.rowptr, mat.storage.col)

    traced_cell = torch.jit.script(my_cell)
    print(traced_cell)
    out = traced_cell(x, adj)
    print(out)
    # # print(traced_cell.code)