import torch


def mult1(x):
    return x.prod(dim=-1).prod(dim=-1)


class Mult(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        y = mult1(x)
        ctx.save_for_backward(x, y)
        return y

    @staticmethod
    def backward(ctx, grad_output):
        x, y = ctx.saved_tensors
        return (grad_output * y)[:, None, None] / x


mult2 = Mult.apply


def check_gradgrad_repeated(x, y):
    gy, = torch.autograd.grad(y[0], x, create_graph=True)
    ggy_1, = torch.autograd.grad(gy[0, 0, 0], x, retain_graph=True)
    gy, = torch.autograd.grad(y[0], x, create_graph=True)
    ggy_2, = torch.autograd.grad(gy[0, 0, 0], x, retain_graph=True)
    print(ggy_1[0, 0, 1], ggy_2[0, 0, 1])


x = torch.ones(2, 4, 4).requires_grad_()
check_gradgrad_repeated(x, mult1(x))
check_gradgrad_repeated(x, mult2(x))