Let’s say f: R^N -> R, and f is batched so that f(xs) -> ys, where xs.shape == (n_batch, n_dim) and ys.shape = (n_batch,). I want to calculate the Laplacian of f, that is, \sum_i d^2/dx_i^2 f(x).

At the moment, I evaluate ys = f(xs), and then iterate over n_dim, and for each i a I calculate

Did you find any solution to this? I also need to calculate Laplacians quite often, my current way of doing this by iterating (see below) is quite slow (and CPU based…):

def laplace(fx: torch.Tensor, x: torch.Tensor):
"""
Laplacian (= sum of 2nd derivations)
of (evaluated) nd->1d-function fx w.r.t. nd-tensor x
:rtype: torch.Tensor
"""
dfx = fx
dfx = torch.autograd.grad(dfx, x, create_graph=True)[0]
ddfx = []
for i in range(len(x)):
vec = torch.tensor([(1 if i == j else 0) for j in range(len(dfx))], dtype=torch.float)
ddfx += [torch.autograd.grad(
dfx,
x,
create_graph=True,
grad_outputs=vec
)[0][i]]
ret = sum(ddfx)
return ret

def laplacian(xs, f, create_graph=False, keep_graph=None, return_grad=False):
xis = [xi.requires_grad_() for xi in xs.flatten(start_dim=1).t()]
xs_flat = torch.stack(xis, dim=1)
ys = f(xs_flat.view_as(xs))
(ys_g, *other) = ys if isinstance(ys, tuple) else (ys, ())
ones = torch.ones_like(ys_g)
(dy_dxs,) = torch.autograd.grad(ys_g, xs_flat, ones, create_graph=True)
lap_ys = sum(
torch.autograd.grad(
dy_dxi, xi, ones, retain_graph=True, create_graph=create_graph
)[0]
for xi, dy_dxi in zip(xis, (dy_dxs[..., i] for i in range(len(xis))))
)
if not (create_graph if keep_graph is None else keep_graph):
ys = (ys_g.detach(), *other) if isinstance(ys, tuple) else ys.detach()
result = lap_ys, ys
if return_grad:
result += (dy_dxs.detach().view_as(xs),)
return result

But actually to calculate the Laplacian, you need to calculate the Hessian, there is no way around that. So as long as your batch is large enough, the for loop should introduce no significant overheard compared to the theoretically optimal implementation. Laplacian will always scale as N^2.