quick update:
for now, I am using the approximation as in this pseudo-code:
def log_softplus(x):
  if x > -5:
    return torch.nn.functional.softplus(x).log()
  else:
    return x

which is numerically stable for small x. If done naively (softplus then log), softplus would return a very small number with low precision.