@albanD
Maybe my description was not very precise. Taking torch.autograd.Function as an example, what I am seeking to do is:
class F(torch.autograd.Function):

   @staticmethod
    def some_func(z1, x2):
         //x2 is also needed here
         return z2

   @staticmethod
    def forward(ctx, x, x2):
         z1 = F.some_func3(x)
         y = z1**2
         z2 = F.some_func(z1, x)

         u = F.some_func2(y) \\derivative wrt x
         ctx.save_for_backward(u)
         return y, z2
                 
    @staticmethod
    def backward(ctx, y.grad, dummygrad):

         u, = ctx.saved_tensors
         return u, None

In this case, I do not need x2’s derivative, but I need to output z2 in the forward pass to save some computation cost in the intermediate values.