Multiple outputs for backward function

Hey,

I am having a hard time to get a working implementation of the following code:

class Linear_custom(Function):

    @staticmethod
    def forward(ctx, input, input_exp, weights, bias):
        
        # Input in integer format covered in float-tensor with additional scaling paramters input_exp
        input_8, input_exp = input, input_exp
      
        weight_8 , exp_w = weights, weights.data_exp
        bias_8, bias_8_exp = bias, bias.data_exp 

        
        ctx.save_for_backward(input_8, weight_8, bias_8)
        ctx.internals = input_exp

    
        output = input_8.mm(weight_8.t()) + bias_8
        
        final_exp = input_exp+exp_w
        out, final_exp = res_shift(output,final_exp)
        ctx.mark_non_differentiable(final_exp)
        return out, final_exp

        

    @staticmethod
    @once_differentiable
    def backward(ctx, grad_output,e_):
        
        
        input_8, weight_8, bias_8 = ctx.saved_tensors
        err_in, err_in_exp = Float_to_int(grad_output)
        weight_exp = weight_8.data_exp
        input_exp = ctx.internals 
        


        grad_input = err_in.mm(weight_8.clone())
        grad_weights = input_8.clone().t().mm(err_in)

        
        grad_bias = grad_output.sum(0)

        err_out, shift_bits = res_shift(grad_input)
        err_exp = err_in_exp + (shift_bits + weight_exp.to(torch.float))

        grad_input = Int_to_float((err_out,err_exp))

       
        
        return grad_input, None, grad_weights.t(), grad_bias
    
        


class Linear(Module):
    def __init__(self,in_features,out_features):
        super(Linear, self).__init__()
        self.fn = Linear_custom.apply
        
        self.weight = nn.Parameter(torch.empty(out_features, in_features))
        self.bias = nn.Parameter(torch.empty(out_features))
        
       
        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))

        fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
        bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0

        nn.init.uniform_(self.bias, -bound, bound)

    def forward(self,input, input_exp):
       
        x = self.fn(input,input_exp,self.weight,self.bias)
        return x

The above code works, and the training seems to work (accuracy & loss )
Now I want to edit the code in the following way:

   @staticmethod
    @once_differentiable
    def backward(ctx, grad_output, grad_output_exp,e_):
        
        
        input_8, weight_8, bias_8 = ctx.saved_tensors
        err_in, err_in_exp = grad_output, grad_output_exp
        weight_exp = weight_8.data_exp
        input_exp = ctx.internals 
        


        grad_input = err_in.mm(weight_8.clone())
        grad_weights = input_8.clone().t().mm(err_in)

        
        grad_bias = grad_output.sum(0)

        err_out, shift_bits = res_shift(grad_input)
        err_exp = err_in_exp + (shift_bits + weight_exp.to(torch.float))

        

       
        
        return err_out, err_exp, None, grad_weights.t(), grad_weight_exp, grad_bias, grad_bias_exp
    

The following new returned values during the backward function call should be (partially) the input for the next backward function. How can I do this ? when returning a tuple(grad_input, grad_input_exp). I will get an error.

return  err_exp, grad_weight_exp, grad_bias_exp

The _exp return variables are Tensors.

Please help :slight_smile:

I don’t know what kind of error you are seeing but would assume it’s pointing to an unexpected number of input/returned values.
The inputs to the backward function correspond to the outputs of the forward and the outputs of the backward correspond to the inputs of the forward. In your example you are increasing the input and output arguments of the backward without manipulating the forward method so a failure is expected since Autograd won’t be able to match these arguments.

Hey, thanks for clarifying this. I essentialy want to return a tuple of (grad_input, grad_input_exponents) as Input for the next backward function call.
Doing this I get following expected error:

torch/autograd/__init__.py", line 154, in backward
    Variable._execution_engine.run_backward(
RuntimeError: expected Variable or None (got tuple)

My question: is there a way to overcome this ?

As previously described, the number of inputs/outputs in the forward/backward has to match. E.g. assuming the forward returns a single tensor while your new backward returns two tensors (grad_input and grad_input_exponents). Where should grad_input_exponents be mapped to given the forward received a single tensor?