Always call jnp.finfo on array dtypes, not arrays themselves (#2147)

* Always call jnp.finfo on array dtypes, not arrays themselves

* Fix instance in safe_normalize

* Remove .lock

Author: dylanhmorris

numpyro/distributions/continuous.py

@@ -348,7 +348,9 @@ def sample(
         assert is_prng_key(key)
         shape = sample_shape + self.batch_shape
         samples = random.dirichlet(key, self.concentration, shape=shape)
-        return jnp.clip(samples, jnp.finfo(samples).tiny, 1 - jnp.finfo(samples).eps)
+        return jnp.clip(
+            samples, jnp.finfo(samples.dtype).tiny, 1 - jnp.finfo(samples.dtype).eps
+        )
 
     @validate_sample
     def log_prob(self, value: ArrayLike) -> ArrayLike:

numpyro/distributions/discrete.py

@@ -1056,7 +1056,7 @@ def entropy(self) -> ArrayLike:
         logq = -jax.nn.softplus(self.logits)
         logp = -jax.nn.softplus(-self.logits)
         p = jax.scipy.special.expit(self.logits)
-        p_clip = jnp.clip(p, jnp.finfo(p).tiny)
+        p_clip = jnp.clip(p, jnp.finfo(p.dtype).tiny)
         return -(1 - p) * logq / p_clip - logp
 
 

numpyro/distributions/util.py

@@ -703,7 +703,7 @@ def safe_normalize(x, *, p=2):
     assert isinstance(p, (float, int))
     assert p >= 0
     norm = jnp.linalg.norm(x, p, axis=-1, keepdims=True)
-    x = x / jnp.clip(norm, jnp.finfo(x).tiny)
+    x = x / jnp.clip(norm, jnp.finfo(x.dtype).tiny)
     # Avoid the singularity.
     mask = jnp.all(x == 0, axis=-1, keepdims=True)
     x = jnp.where(mask, x.shape[-1] ** (-1 / p), x)