From 2669e4cf843bdde64bb95d995fda374313db9deb Mon Sep 17 00:00:00 2001 From: = Date: Sun, 2 Jun 2024 19:30:58 -0400 Subject: [PATCH] example of bug --- bug.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bug.py b/bug.py index 713ebd9c9..cd3a95285 100644 --- a/bug.py +++ b/bug.py @@ -103,7 +103,7 @@ def run_mclmc(logdensity_fn, num_steps, initial_position, transform, key, precon key = jax.random.PRNGKey(0) -map = jax.pmap # change to pmap +map = jax.pmap # change to vmap sampler = partial(run_mclmc, preconditioning=True)