Heterogeneous message passing

[HamidL]

Hi everyone!
I want to use a heterogeneous graph to perform node classification of a certain node type with PyG. Let me explain the task:

Heterogeneous graph

• Node type: connection and host
• Edge types: (connection, from, host) and (host, to, connection)

My HeteroData object is the following:

HeteroData(
  connection={
    num_nodes=200,
    x=[200, 26],
    y=[200, 15]
  },
  host={
    num_nodes=2,
    x=[2, 9]
  },
  (connection, from, host)={ edge_index=[2, 400] },
  (host, to, connection)={ edge_index=[2, 400] }
)

The task

1. Classify each of the connection nodes (one-hot encoding of dimension 15).

Desired model

I've started to implement a model, but I don't know how to do a few things. The current model is:

class GNN(MessagePassing):
    def __init__(self, hidden_channels=128):

        super().__init__(aggr="mean")

        self.message_func_host = nn.Sequential(
            nn.Dropout(p=0.5),
            nn.Linear(hidden_channels, hidden_channels)
        )
        self.message_func_conn = nn.Sequential(
            nn.Dropout(p=0.5),
            nn.Linear(hidden_channels, hidden_channels)
        )

        self.host_update = nn.GRU(hidden_channels, hidden_channels)
        self.connection_update = nn.GRU(hidden_channels, hidden_channels)

        self.readout_nn = nn.Sequential(
            nn.Linear(hidden_channels, hidden_channels),
            nn.ReLU(),
            nn.Dropout(p=0.5),
            nn.Linear(hidden_channels, 64),
            nn.ReLU(),
            nn.Dropout(p=0.5),
            nn.Linear(64, 15),
            nn.Softmax()
        )

    def forward(self, x: Tensor, edge_index: Adj):
        # TODO: implement

        return self.propagate(edge_index, x=x)

    def readout(self, states):
        return self.readout_nn(states)

    def update_host(self, x):
        return self.host_update(x)

    def update_connection(self, x):
        return self.connection_update(x)

My questions are:

1. As the names states, I'd like to use host_update and connection_update for updating each of these types (even though now they are the same unit, I'd like to change them in the future, which is why I'm creating two).
2. The same happens with message_func_host and message_func_conn. I'd like to use two equal functions to have different weights for each node type.
3. I'd like to initialize the hidden states of the nodes in the following way: for the host nodes, use a 128-length tensor of 0's; for the connection nodes, concatenate the 26-length tensor of attributes data["connection"].x with a 102-length tensor of 0's.
4. I'd like to implement the message passing process T times for every graph. Would I need to create a loop in the forward function?? I don't know what would be the best way to implement something like this.

By reading the docs, I understand that step 2 could be achieved by using HeteroConv, but could the update function be implemented with it? And if so, how?

If more details are needed, please just ask!
Thank you very much for your time.