Format fix

zaxtax · zaxtax · commit f5aeaf68c683 · 2021-11-17T16:00:21.000Z
diff --git a/pymc/sampling_jax.py b/pymc/sampling_jax.py
@@ -173,10 +173,10 @@ def sample_numpyro_nuts(
     map_seed = jax.random.split(seed, chains)
 
     if chains == 1:
-        init_params=init_state
+        init_params = init_state
         map_seed = seed
     else:
-        init_params=init_state_batched
+        init_params = init_state_batched
 
     pmap_numpyro.run(
         map_seed,
