Adding coords, dims, and some optimizations

Author: zaxtax

pymc/sampling_jax.py

@@ -124,8 +124,9 @@ def _get_log_likelihood(model, samples):
     for v in model.observed_RVs:
         logp_v = replace_shared_variables([logpt(v)])
         fgraph = FunctionGraph(model.value_vars, logp_v, clone=False)
+        optimize_graph(fgraph, include=["fast_run"], exclude=["cxx_only", "BlasOpt"])
         jax_fn = jax_funcify(fgraph)
-        result = jax.vmap(jax.vmap(jax_fn))(*samples)[0]
+        result = jax.jit(jax.vmap(jax.vmap(jax_fn)))(*samples)[0]
         data[v.name] = result
     return data
 
@@ -150,6 +151,20 @@ def sample_numpyro_nuts(
 
     vars_to_sample = list(get_default_varnames(var_names, include_transformed=keep_untransformed))
 
+    coords = {
+        cname: np.array(cvals) if isinstance(cvals, tuple) else cvals
+        for cname, cvals in model.coords.items()
+        if cvals is not None
+    }
+
+    if hasattr(model, "RV_dims"):
+        dims = {
+            var_name: [dim for dim in dims if dim is not None]
+            for var_name, dims in model.RV_dims.items()
+        }
+    else:
+        dims = {}
+
     tic1 = pd.Timestamp.now()
     print("Compiling...", file=sys.stdout)
 
@@ -213,6 +228,7 @@ def sample_numpyro_nuts(
     mcmc_samples = {}
     for v in vars_to_sample:
         fgraph = FunctionGraph(model.value_vars, [v], clone=False)
+        optimize_graph(fgraph, include=["fast_run"], exclude=["cxx_only", "BlasOpt"])
         jax_fn = jax_funcify(fgraph)
         result = jax.vmap(jax.vmap(jax_fn))(*raw_mcmc_samples)[0]
         mcmc_samples[v.name] = result
@@ -221,11 +237,13 @@ def sample_numpyro_nuts(
     print("Transformation time = ", tic4 - tic3, file=sys.stdout)
 
     posterior = mcmc_samples
-    az_posterior = az.from_dict(posterior=posterior)
-
-    az_obs = az.from_dict(observed_data=find_observations(model))
-    az_stats = az.from_dict(sample_stats=_sample_stats_to_xarray(pmap_numpyro))
-    az_ll = az.from_dict(log_likelihood=_get_log_likelihood(model, raw_mcmc_samples))
-    az_trace = az.concat(az_posterior, az_ll, az_obs, az_stats)
+    az_trace = az.from_dict(
+        posterior=posterior,
+        log_likelihood=_get_log_likelihood(model, raw_mcmc_samples),
+        observed_data=find_observations(model),
+        sample_stats=_sample_stats_to_xarray(pmap_numpyro),
+        coords=coords,
+        dims=dims,
+    )
 
     return az_trace