Hi,
You can do the following:
class Linear(Function):
    @staticmethod
    def forward(ctx, my_custom_logic, input, weight, bias=None):
        ctx.my_custom_logic = my_custom_logic
        # ...
        # Compute output taking advantage of ctx.my_custom_logic
        # ...
        return output

    @staticmethod
    def backward(ctx, grad_output):
        grad_input = grad_weight = grad_bias = None
        # ...
        # Compute grads taking advantage of ctx.my_custom_logic
        # ...
        return grad_input, grad_weight, grad_bias