How to calculate geometric mean in a differentiable way?

Question

How to calculate goemetric mean along a dimension using Pytorch? Some numbers can be negative. The function must be differentiable.

Answer 1

A known (reasonably) numerically-stable version of the geometric mean is:

import torch

def gmean(input_x, dim):
    log_x = torch.log(input_x)
    return torch.exp(torch.mean(log_x, dim=dim))

x = torch.Tensor([2.0] * 1000).requires_grad_(True)
print(gmean(x, dim=0))
# tensor(2.0000, grad_fn=<ExpBackward>)

This kind of implementation can be found, for example, in SciPy (see here), which is a quite stable lib.

---

The implementation above does not handle zeros and negative numbers. Some will argue that the geometric mean with negative numbers is not well-defined, at least when not all of them are negative.