def METHOD_NAME(self, value, keepdim=False):
    weights = self._log_weights.reshape(
        self._log_weights.size()
        + torch.Size([1] * (value.dim() - self._log_weights.dim()))
    )
    dim = self._aggregation_dim
    max_weight = weights.max(dim=dim, keepdim=True)[0]
    relative_probs = (weights - max_weight).exp()
    return (value * relative_probs).sum(
        dim=dim, keepdim=keepdim
    ) / relative_probs.sum(dim=dim, keepdim=keepdim)