import torch


def geometric_mean(x: torch.Tensor, dim: int) -> torch.Tensor:
    """
    Compute geometric mean along given dimension.
    This implementation assume values are in range (0...1) (Probabilities)
    Args:
        x: Input tensor of arbitrary shape
        dim: Dimension to reduce
    Returns:
        Tensor
    """
    return x.log().mean(dim=dim).exp()