mlpg.py 文件源码

def backward(self, grad_output):
        means, = self.saved_tensors
        T = self.R.shape[0]
        dim = means.dim()

        # Add batch axis if necessary
        if dim == 2:
            T_, D = means.shape
            B = 1
            grad_output = grad_output.view(B, T, -1)
        else:
            B, T_, D = means.shape

        grad = torch.matmul(self.R.transpose(0, 1), grad_output)

        reshaped = not (T == T_)
        if not reshaped:
            grad = grad.view(B, self.num_windows, T, -1).transpose(
                1, 2).contiguous().view(B, T, D)

        if dim == 2:
            return grad.view(-1, D)

        return grad