How to use HeteroBatchNorm and HeteroLayerNorm?

[RENEK-bool]

I noticed that HeteroBatchNorm requires x and type_vec as inputs, and HeteroLayerNorm requires x, type_vec and type_ptr as inputs.

1. How should type_vec and type_ptr be constructed?
2. Do I need to call HeteroBatchNorm twice to process nodes and edges respectively?
   My heterogeneous graph has one type of nodes and two types of edges.
   Below is my code.

class GNN(torch.nn.Module):
    def __init__(self, hidden_channels):
        super().__init__()
        
        self.wire_convs = torch.nn.ModuleList()
        for _ in range(3):
            conv = HeteroConv({
                ('pin', 'cell_arc', 'pin'): TransformerConv((-1, -1), hidden_channels, edge_dim=-1),
                ('pin', 'net_arc', 'pin'): TransformerConv((-1, -1), hidden_channels, edge_dim=-1),
            }, aggr='sum')
            self.wire_convs.append(conv)

        self.act = nn.ReLU()
        self.wire_fc1 = nn.Linear(hidden_channels,32)
        self.wire_fc2 = nn.Linear(32,16)
        self.wire_fc3 = nn.Linear(16,11)

    def forward(self, x_dict, edge_index_dict, edge_attr_dict):
        for wire_conv in self.wire_convs:
            x_dict = wire_conv(x_dict, edge_index_dict, edge_attr_dict)
            x_dict = {key: x.relu() for key, x in x_dict.items()}
        
        src, dst = edge_index_dict['pin','net_arc','pin']
        edge_attr = torch.cat([x_dict['pin'][src], x_dict['pin'][dst]], dim=-1)

        x = self.wire_fc1(edge_attr)
        x = self.act(x)
        x = self.wire_fc2(x)
        x = self.act(x)
        x = self.wire_fc3(x)
    
        return x

[rusty1s]

How should type_vec and type_ptr be constructed?

This refers to the same data just in a different format. type_vec refers to the classic [0, 0, 0, 1, 1, 1, 2, 2, 2, ...] index representation which assigns each node to a specific node type. type_ptr refers to its compressed representation, i.e., it only stores its boundaries, e.g.: [0, 4, 7] for two node types where node type 1 has 4 nodes and node type 2 has 3 nodes.

If you are using dictionaries to process heterogeneous graphs, you can just do

norms = torch.nn.ModuleDict({node_type: BatchNorm(...) for node_type in node_types})

x_dict = {key: norms[key](x) for x, key in x_dict.items()}