Modify trainable affinity matrix (parameter) to calculate Laplacian Matrix

Hi there!

I have a problem in which part of the loss function depends on a trainable affinity matrix W and pairwise differences between elements of the output, which looks like this: \sum_{ij} w_ij * (y_i - y_j), where w_ij represents the affinity between position i and j. The previous equation can be written in terms of the Laplacian matrix: y*L*y^T where L = D - W and D, the degree matrix, is a diagonal matrix so that D_ii = \sum_j w_ij (each element of the diagonal is the sum of the elements of the corresponding row in the affinity matrix).

The thing is that I want to learn W but it is easier to use L instead, so I created a module to calculate my loss that would look like this:

class CRFlayer(nn.Module, ABC):
    def __init__(self, nodes=25):
        super(CRFlayer, self).__init__()
        self.nodes = nodes
        self.W = nn.Parameter(zeros(nodes, nodes))
        self.W.requires_grad = True

    def forward(self, Y, H):
        # Calculate Degree matrix using W
        D = "???"
        # Calculate Laplacian matrix using W
        L = D - self.W
        # Calculate loss based on L, the output Y, and other inputs
        loss = ...
        return loss

My question is that if it’s possible to do that. In particular, I don’t know how to code the degree matrix D and I don’t know if the gradients will be correctly calculated over W after that. Do you have any recommendations?