From cc1ad02c5aab0586a5a613ef6f2516d5f8848b10 Mon Sep 17 00:00:00 2001 From: Frans Zdyb Date: Thu, 17 Aug 2023 22:48:51 +0200 Subject: [PATCH] Fix bug in SineBivariateVonMises sampler (#1628) * sbvm sampler fix by comparison with R code * use rsqrt instead of 1/sqrt --- numpyro/distributions/directional.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/numpyro/distributions/directional.py b/numpyro/distributions/directional.py index b3e1dcd1a..67a91a8d1 100644 --- a/numpyro/distributions/directional.py +++ b/numpyro/distributions/directional.py @@ -477,7 +477,7 @@ def update_fn(curr): phi_key, key = random.split(key) accept_key, acg_key, phi_key = random.split(phi_key, 3) - x = jnp.sqrt(1 + 2 * eig / b0) * random.normal(acg_key, shape) + x = lax.rsqrt(1 + 2 * eig / b0) * random.normal(acg_key, shape) x /= jnp.linalg.norm( x, axis=1, keepdims=True ) # Angular Central Gaussian distribution