I’m trying to backprop through a higher-order function (a function that takes a function as argument), specifically a functional 1 (a higher-order function that returns a scalar). Here is a simple example:
import torch

class Functional(torch.autograd.Function):
    @staticmethod
    def forward(ctx, f):
        value = f(2)**2 - f(1)
        ctx.save_for_backward(value)
        return value
    @staticmethod
    def backward(ctx, grad_output):
        value = ctx.saved_tensors
        value.backward(grad_output)
        return None

class Function(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.a = torch.tensor(0.)
    def forward(self, x):
        return self.a * x

function = Function()
functional = Functional.apply

value = functional(function)
value.backward() # RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

Unfortunately, this throws an error. How can I fix this?
Here is another example, a Monte Carlo estimate of the inner product between two functions on the unit hypercube:
import torch

dim = 2
samples = 100

class InnerProduct(torch.autograd.Function):
    @staticmethod
    def forward(ctx, f, g):
        x = torch.rand(samples, dim)
        f_x = f(x)
        g_x = g(x)
        ctx.save_for_backward(f_x, g_x)
        return (f_x * g_x).mean(-1)
    @staticmethod
    def backward(ctx, grad_output):
        f_x, g_x = ctx.saved_tensors
        for f_x_i, g_x_i in zip(f_x, g_x):
            f_x_i.backward(g_x_i / samples * grad_output)
            g_x_i.backward(f_x_i / samples * grad_output)
        return None, None
inner_product = InnerProduct.apply

class ScalarFunction(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.net = torch.nn.Sequential(
            torch.nn.Linear(dim, 10),
            torch.nn.ELU(),
            torch.nn.Linear(10, 10),
            torch.nn.ELU(),
            torch.nn.Linear(10, 1),
            torch.nn.Flatten(-2, -1)
        )
    def forward(self, x):
        return self.net(x)

f = ScalarFunction()
g = ScalarFunction()
a = inner_product(f, g)
print(a)
a.backward() # RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn