log1mexp vmap fix

Author: jjcmoon

experiments/visual_sudoku/run.py

@@ -96,7 +96,7 @@ def get_circuit(grid_size: int):
 
 
 def nll_loss(preds, targets):
-    neg_preds = klay.torch.log1mexp(preds, eps=1e-7)
+    neg_preds = klay.torch.utils.log1mexp(preds)
     nll = -torch.where(targets, preds, neg_preds)
     return nll.mean()
 

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
+    mask = CUTOFF < x
+    return torch.where(
+        mask,
+        (-x.clamp(min=CUTOFF).expm1() + eps).log(),
+        (-x.clamp(max=CUTOFF).exp() + eps).log1p()
+    )
 
 
 def negate_real(x, eps):