Message Passing Layer of PointGNN

[andreasbinder]

I closed this pull request by mistake and solved the conflicts

[codecov]

Codecov Report

Merging #6194 (ffc1d07) into master (90de72f) will increase coverage by 0.02%.
The diff coverage is 100.00%.

❗ Current head ffc1d07 differs from pull request most recent head c2c2ccb. Consider uploading reports for the commit c2c2ccb to get more accurate results

@@            Coverage Diff             @@
##           master    #6194      +/-   ##
==========================================
+ Coverage   84.51%   84.53%   +0.02%     
==========================================
  Files         372      373       +1     
  Lines       20805    20833      +28     
==========================================
+ Hits        17584    17612      +28     
  Misses       3221     3221              

Impacted Files	Coverage Δ
torch_geometric/nn/conv/__init__.py	100.00% <100.00%> (ø)
torch_geometric/nn/conv/point_gnn_conv.py	100.00% <100.00%> (ø)

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[rusty1s]

I think this looks great. Can we also add a simple test for this?

[andreasbinder]

Hi @rusty1s ! I am sorry to bother you but I am not sure how to write the jittable part of the test (or am I even supposed to run that part, I removed the if clause in my local test).

x = torch.randn(6, 3)
edge_index = torch.tensor([[0, 1, 1, 1, 2, 5], [1, 2, 3, 4, 3, 4]])
row, col = edge_index
adj = SparseTensor(row=row, col=col, sparse_sizes=(6, 6))
MLP_h = MLP([3, 64, 3])
MLP_f = MLP([6, 64, 3])
MLP_g = MLP([3, 64, 3])

conv = PointGNNConv(state_channels=3, MLP_h=MLP_h, MLP_f=MLP_f, MLP_g=MLP_g)
assert conv.__repr__() == 'PointGNNConv(3, lin_h=MLP(3, 64, 3), lin_f=MLP(6, 64, 3),lin_g=MLP(3, 64, 3))'
out = conv(x, x, edge_index)

assert out.size() == (6, 3)
assert torch.allclose(conv(x, x, adj.t()), out, atol=1e-6)

if True:
    t = '(Tensor, Tensor, Tensor) -> Tensor'
    jit = torch.jit.script(conv.jittable(t))
    assert jit(x, x, edge_index).tolist() == out.tolist()

    t = '(Tensor, Tensor, SparseTensor) -> Tensor'
    jit = torch.jit.script(conv.jittable(t))
    assert torch.allclose(jit(x, x, adj.t()), out, atol=1e-5)

I could not infer the error from the trace:

I can attach the whole trace if you think it helps. Thank you very much!

[rusty1s]

Don't worry about it, I can take care of it. Do you want to push your tests so that I can take a look?

[andreasbinder]

Amazing! I added the test file (and some annotations to the conv file) @rusty1s

[andreasbinder]

Thank you for the useful changes @rusty1s ! Just out of interest, do you know why the jittable test failed (so that I know in the future)? Thank you!

[rusty1s]

I am not so sure anymore. I think it had something do to with the missing size=None in propagate which we sadly currently need for TorchScript support :(