I’m defining a new function using the 0.2 style and am wondering when it is appropriate to store intermediate results in the ctx object as opposed to using the save_for_backward function.
Here is a simple example:
# OPTION 1
class Square(Function):
    @staticmethod
    def forward(ctx, a, b):
        ctx.save_for_backward(a, b)
        c = a + b
        return c * c

    @staticmethod
    def backward(ctx, grad_output):
        a, b = ctx.saved_variables
        c = a + b

        grad = grad_output * (2 * c)
        return grad, grad


# OPTION 2
class Square(Function):
    @staticmethod
    def forward(ctx, a, b):
        c = a + b
        ctx.intermediate_results = c
        return c * c

    @staticmethod
    def backward(ctx, grad_output):
        c = ctx.intermediate_results

        grad = grad_output * (2 * c)
        return grad, grad


What are the advantages/disadvantages to Option 1 over Option 2?  When would you want to use one over the other?
(In a real example, the computation of the intermediate variable “c” would be a lot more intensive, e.g. involve O(n^2) operations)
Thank you.