log1mexp vmap fix

Author: jjcmoon

src/klay/torch/utils.py

@@ -5,17 +5,18 @@
 CUTOFF = -math.log(2)
 
 
-def log1mexp(x, eps=10e-12):
+def log1mexp(x, eps=1e-12):
     """
     Numerically accurate evaluation of log(1 - exp(x)) for x < 0.
     See [Maechler2012accurate]_ for details.
     https://github.com/pytorch/pytorch/issues/39242
     """
     mask = CUTOFF < x  # x < 0
-    out = torch.empty_like(x)
-    out[mask] = (-x[mask].expm1() + eps).log()
-    out[~mask] = (-x[~mask].exp() + eps).log1p()
-    return out
+    return torch.where(
+        mask,
+        (-x.clamp(min=CUTOFF).expm1() + eps).log(),
+        (-x.clamp(max=CUTOFF).exp() + eps).log1p()
+    )
 
 
 def negate_real(x, eps):