Message Passing with multiple Tensors in TorchScript

[dav0dea]

Hi, I implemented a message passing layer based on MessagePassing, which returns a tuple with two Tensors from message. I then aggregate both of them by overriding the aggregate function. All of this works as intended.
However, when I use TorchScript with this model I run into some typing problems. I am creating a new update function with proper type annotations like following to get rid of the typing problems there:

def update(self, inputs: Tuple[torch.Tensor, torch.Tensor]) -> Tuple[torch.Tensor, torch.Tensor]:
    return inputs

After that I get the following error in the generated propagate function:

Return value was annotated as having type Tensor but is actually of type Tuple[Tensor, Tensor]:
  File "/tmp/usr_pyg_jit/tmpzaho5dck.py", line 225
        out = self.message(q_i=kwargs.q_i, dk=kwargs.dk, d_ij=kwargs.d_ij, r_ij=kwargs.r_ij, vec_j=kwargs.vec_j, k_j=kwargs.k_j, v_j=kwargs.v_j, dv=kwargs.dv)
        out = self.aggregate(out, dim_size=kwargs.dim_size, ptr=kwargs.ptr, index=kwargs.index)
        return self.update(out)
        ~~~~~~~~~~~~~~~~~~~~~~ <--- HERE

I believe the return type annotations for propagate are automatically generated to be Tensor when calling .jittable(). Is there a way to annotate the propagate return type, maybe through the propagate_type annotation that is already present?
If not or if this is related to something else, is the MessagePassing module intended to support messages made up of multiple tensors or do you have other suggestions to achieve the same thing?

[rusty1s]

Thanks for reporting. I made some changes and installing from master should fix it, see here.

[dav0dea]

Thank you for looking into this! However, the new test failed on the most recent commit.

[rusty1s]

Yeah, there were some bugs in Python 3.6 which are now resolved.