FunctionErrorFeedback.py

import torch
from torch.autograd import Function, Variable


class ErrorFeedbackFunction(Function):
    # the forward pass consists in copying the input to the output
    @staticmethod
    def forward(ctx, input, feedback):
        ctx.save_for_backward(input, feedback)
        return input
    
    @staticmethod
    def backward(ctx, grad_output):
    # grad_feedback is None and grad_input is the random projection
    # of the error
        input, feedback = ctx.saved_variables
        grad_input = grad_feedback = None

        if ctx.needs_input_grad[0]:
            # random projection
            grad_input = torch.mm(grad_output, feedback)

        return grad_input, grad_feedback