Thanks for questioning my code because I think it was wrong…if nothing else it was confusing enough that I couldn’t understand it. I decided to structure my data slightly different so the code could access it a bit nicer.
Now it looks like this:


class NeighborhoodEmbedding(nn.Module):
    """Combines node embeddings for a node and its neighbors/children into a single vector."""

    def forward(self, fs: torch.Tensor, tree: List[List[int]]):
        childrenVecs = [fs[t] for t in tree]

        # average each node's representation with its children's representations as vectors.
        vecs = [torch.cat([x.reshape(1,-1),y]).mean(0) for x,y in zip(fs, childrenVecs)]
        return torch.stack(vecs)


class TreeEmbedding(nn.Module):
    """Turns a clause/context tree into a vector."""
    
    def __init__(self, signatureSize, symbolDim):
        super().__init__()
        self.symbolVectors = torch.nn.Parameter(torch.randn(signatureSize, symbolDim))

        self.nodeEmbedding = NodeEmbedding(symbolDim, symbolDim)
        self.neighborhoodEmbedding = NeighborhoodEmbedding()

    def forward(self, tree: List[List[int]], 
                      treeSymbols: List[int]):
        """tree lists the children of each node. 
        So if tree[3] == [6,10], this means that node 3 has two children 
        and that its grandchildren are specified at tree[6] and tree[10]"""

        nodeVecs = self.symbolVectors[treeSymbols]

        for i in range(getTreeDepth(tree)):
            fs = self.nodeEmbedding(nodeVecs)
            nodeVecs = self.neighborhoodEmbedding(fs, tree)

        out = nodeVecs[0]
        return out

I think this works and is a little more efficient than before since that triply nested for-loop is now only doubly nested with a pytorch operation kind of working as a stand-in for the other loop.
I think I could do it in a better way if pytorch supported. For instance I wish [fs[t] for t in tree] could be simply fs[tree]
Can you see any optimizations I can perform to make this better?
Is it now clearer what computation I’m performing?
The self.nodeEmbedding is simply a shallow fully connected NN which doesn’t change the dimension of its input. The self.neighborEmbedding is meant to take the transformed node vectors and the tree (explained in the docstring in the code) and return updated node vectors. Right now I’m implementing this simply as each node’s vector getting mapped to the average of it and its children.
Thanks again for all your help!