improve precision for constraint's logdet (#158)

numpyro/distributions/constraints.py

@@ -331,7 +331,8 @@ def inv(self, y):
         return logit(y)
 
     def log_abs_det_jacobian(self, x, y):
-        return np.log(y * (1 - y))
+        x_abs = np.abs(x)
+        return -x_abs - 2 * np.log1p(np.exp(-x_abs))
 
 
 class StickBreakingTransform(Transform):
@@ -361,8 +362,11 @@ def inv(self, y):
     def log_abs_det_jacobian(self, x, y):
         # Ref: https://mc-stan.org/docs/2_19/reference-manual/simplex-transform-section.html
         # |det|(J) = Product(y * (1 - z))
-        z = _clipped_expit(x - np.log(x.shape[-1] - np.arange(x.shape[-1])))
-        return np.sum(np.log(y[..., :-1] * (1 - z)), axis=-1)
+        x = x - np.log(x.shape[-1] - np.arange(x.shape[-1]))
+        z = np.clip(expit(x), a_min=np.finfo(x.dtype).tiny)
+        # XXX we use the identity 1 - z = z * exp(-x) to not worry about
+        # the case z ~ 1
+        return np.sum(np.log(y[..., :-1] * z) - x, axis=-1)
 
 
 ##########################################################