Transform samples from sample_numpyro_nuts

* Add `pymc3.sampling_jax._transform_samples` function which transforms draws

* Modify `pymc3.sampling_jax.sample_numpyro_nuts` function to use this function to return transformed samples

* Add release note

Author: martiningram

Diff for: RELEASE-NOTES.md

@@ -5,6 +5,7 @@
 ### Breaking Changes
 
 ### New Features
+- `pymc3.sampling_jax.sample_numpyro_nuts` now returns samples from transformed random variables, rather than from the unconstrained representation (see [#4427](https://github.com/pymc-devs/pymc3/pull/4427)).
 
 ### Maintenance
 - `math.log1mexp_numpy` no longer raises RuntimeWarning when given very small inputs. These were commonly observed during NUTS sampling (see [#4428](https://github.com/pymc-devs/pymc3/pull/4428)).

Diff for: pymc3/sampling_jax.py

@@ -3,6 +3,8 @@
 import re
 import warnings
 
+from collections import defaultdict
+
 xla_flags = os.getenv("XLA_FLAGS", "").lstrip("--")
 xla_flags = re.sub(r"xla_force_host_platform_device_count=.+\s", "", xla_flags).split()
 os.environ["XLA_FLAGS"] = " ".join(["--xla_force_host_platform_device_count={}".format(100)])
@@ -175,8 +177,50 @@ def _sample(current_state, seed):
     # print("Sampling time = ", tic4 - tic3)
 
     posterior = {k: v for k, v in zip(rv_names, mcmc_samples)}
+    tic3 = pd.Timestamp.now()
+    posterior = _transform_samples(posterior, model, keep_untransformed=False)
+    tic4 = pd.Timestamp.now()
 
     az_trace = az.from_dict(posterior=posterior)
-    tic3 = pd.Timestamp.now()
     print("Compilation + sampling time = ", tic3 - tic2)
+    print("Transformation time = ", tic4 - tic3)
+
     return az_trace  # , leapfrogs_taken, tic3 - tic2
+
+
+def _transform_samples(samples, model, keep_untransformed=False):
+
+    # Find out which RVs we need to compute:
+    free_rv_names = {x.name for x in model.free_RVs}
+    unobserved_names = {x.name for x in model.unobserved_RVs}
+
+    names_to_compute = unobserved_names - free_rv_names
+    ops_to_compute = [x for x in model.unobserved_RVs if x.name in names_to_compute]
+
+    # Create function graph for these:
+    fgraph = theano.graph.fg.FunctionGraph(model.free_RVs, ops_to_compute)
+
+    # Jaxify, which returns a list of functions, one for each op
+    jax_fns = jax_funcify(fgraph)
+
+    # Put together the inputs
+    inputs = [samples[x.name] for x in model.free_RVs]
+
+    for cur_op, cur_jax_fn in zip(ops_to_compute, jax_fns):
+
+        # We need a function taking a single argument to run vmap, while the
+        # jax_fn takes a list, so:
+        to_run = lambda x: cur_jax_fn(*x)
+
+        result = jax.vmap(jax.vmap(to_run))(inputs)
+
+        # Add to sample dict
+        samples[cur_op.name] = result
+
+    # Discard unwanted transformed variables, if desired:
+    vars_to_keep = set(
+        pm.util.get_default_varnames(list(samples.keys()), include_transformed=keep_untransformed)
+    )
+    samples = {x: y for x, y in samples.items() if x in vars_to_keep}
+
+    return samples