Added a small test

Author: martiningram

Diff for: pymc3/sampling_jax.py

@@ -123,6 +123,7 @@ def sample_numpyro_nuts(
     random_seed=10,
     model=None,
     progress_bar=True,
+    keep_untransformed=False,
 ):
     from numpyro.infer import MCMC, NUTS
 
@@ -178,7 +179,7 @@ def _sample(current_state, seed):
 
     posterior = {k: v for k, v in zip(rv_names, mcmc_samples)}
     tic3 = pd.Timestamp.now()
-    posterior = _transform_samples(posterior, model, keep_untransformed=False)
+    posterior = _transform_samples(posterior, model, keep_untransformed=keep_untransformed)
     tic4 = pd.Timestamp.now()
 
     az_trace = az.from_dict(posterior=posterior)

Diff for: pymc3/tests/test_sampling_jax.py

@@ -0,0 +1,19 @@
+import numpy as np
+
+import pymc3 as pm
+
+from pymc3.sampling_jax import sample_numpyro_nuts
+
+
+def test_transform_samples():
+
+    with pm.Model() as model:
+
+        sigma = pm.HalfNormal("sigma")
+        b = pm.Normal("b", sigma=sigma)
+        trace = sample_numpyro_nuts(keep_untransformed=True)
+
+    log_vals = trace.posterior["sigma_log__"].values
+    trans_vals = trace.posterior["sigma"].values
+
+    assert np.allclose(np.exp(log_vals), trans_vals)