refactor sampling_jax postrocessing to avoid jit (#5908)

Refactor sampling_jax postrocessing to avoid JIT and multiple function definitions

Co-authored-by: Ricardo Vieira <[email protected]>

Author: ferrine

pymc/sampling_jax.py

@@ -134,13 +134,10 @@ def _sample_stats_to_xarray(posterior):
 
 def _get_log_likelihood(model: Model, samples, backend=None) -> Dict:
     """Compute log-likelihood for all observations"""
-    data = {}
-    for v in model.observed_RVs:
-        v_elemwise_logp = model.logp(v, sum=False)
-        jax_fn = get_jaxified_graph(inputs=model.value_vars, outputs=v_elemwise_logp)
-        result = jax.jit(jax.vmap(jax.vmap(jax_fn)), backend=backend)(*samples)[0]
-        data[v.name] = result
-    return data
+    elemwise_logp = model.logp(model.observed_RVs, sum=False)
+    jax_fn = get_jaxified_graph(inputs=model.value_vars, outputs=elemwise_logp)
+    result = jax.vmap(jax.vmap(jax_fn))(*jax.device_put(samples, jax.devices(backend)[0]))
+    return {v.name: r for v, r in zip(model.observed_RVs, result)}
 
 
 def _get_batched_jittered_initial_points(
@@ -339,13 +336,11 @@ def sample_blackjax_nuts(
     print("Sampling time = ", tic3 - tic2, file=sys.stdout)
 
     print("Transforming variables...", file=sys.stdout)
-    mcmc_samples = {}
-    for v in vars_to_sample:
-        jax_fn = get_jaxified_graph(inputs=model.value_vars, outputs=[v])
-        result = jax.jit(jax.vmap(jax.vmap(jax_fn)), backend=postprocessing_backend)(
-            *raw_mcmc_samples
-        )[0]
-        mcmc_samples[v.name] = result
+    jax_fn = get_jaxified_graph(inputs=model.value_vars, outputs=vars_to_sample)
+    result = jax.vmap(jax.vmap(jax_fn))(
+        *jax.device_put(raw_mcmc_samples, jax.devices(postprocessing_backend)[0])
+    )
+    mcmc_samples = {v.name: r for v, r in zip(vars_to_sample, result)}
 
     tic4 = datetime.now()
     print("Transformation time = ", tic4 - tic3, file=sys.stdout)
@@ -355,10 +350,14 @@ def sample_blackjax_nuts(
     else:
         idata_kwargs = idata_kwargs.copy()
 
-    if idata_kwargs.pop("log_likelihood", True):
+    if idata_kwargs.pop("log_likelihood", bool(model.observed_RVs)):
+        tic5 = datetime.now()
+        print("Computing Log Likelihood...", file=sys.stdout)
         log_likelihood = _get_log_likelihood(
             model, raw_mcmc_samples, backend=postprocessing_backend
         )
+        tic6 = datetime.now()
+        print("Log Likelihood time = ", tic6 - tic5, file=sys.stdout)
     else:
         log_likelihood = None
 
@@ -531,13 +530,11 @@ def sample_numpyro_nuts(
     print("Sampling time = ", tic3 - tic2, file=sys.stdout)
 
     print("Transforming variables...", file=sys.stdout)
-    mcmc_samples = {}
-    for v in vars_to_sample:
-        jax_fn = get_jaxified_graph(inputs=model.value_vars, outputs=[v])
-        result = jax.jit(jax.vmap(jax.vmap(jax_fn)), backend=postprocessing_backend)(
-            *raw_mcmc_samples
-        )[0]
-        mcmc_samples[v.name] = result
+    jax_fn = get_jaxified_graph(inputs=model.value_vars, outputs=vars_to_sample)
+    result = jax.vmap(jax.vmap(jax_fn))(
+        *jax.device_put(raw_mcmc_samples, jax.devices(postprocessing_backend)[0])
+    )
+    mcmc_samples = {v.name: r for v, r in zip(vars_to_sample, result)}
 
     tic4 = datetime.now()
     print("Transformation time = ", tic4 - tic3, file=sys.stdout)
@@ -547,10 +544,14 @@ def sample_numpyro_nuts(
     else:
         idata_kwargs = idata_kwargs.copy()
 
-    if idata_kwargs.pop("log_likelihood", True):
+    if idata_kwargs.pop("log_likelihood", bool(model.observed_RVs)):
+        tic5 = datetime.now()
+        print("Computing Log Likelihood...", file=sys.stdout)
         log_likelihood = _get_log_likelihood(
             model, raw_mcmc_samples, backend=postprocessing_backend
         )
+        tic6 = datetime.now()
+        print("Log Likelihood time = ", tic6 - tic5, file=sys.stdout)
     else:
         log_likelihood = None
 

pymc/tests/test_sampling_jax.py

@@ -1,5 +1,6 @@
 import aesara
 import aesara.tensor as at
+import jax
 import numpy as np
 import pytest
 
@@ -27,7 +28,16 @@
     ],
 )
 @pytest.mark.parametrize("postprocessing_backend", [None, "cpu"])
-def test_transform_samples(sampler, postprocessing_backend):
+@pytest.mark.parametrize(
+    "chains",
+    [
+        pytest.param(1),
+        pytest.param(
+            2, marks=pytest.mark.skipif(len(jax.devices()) < 2, reason="not enough devices")
+        ),
+    ],
+)
+def test_transform_samples(sampler, postprocessing_backend, chains):
     aesara.config.on_opt_error = "raise"
     np.random.seed(13244)
 
@@ -39,7 +49,7 @@ def test_transform_samples(sampler, postprocessing_backend):
         b = pm.Normal("b", a, sigma=sigma, observed=obs_at)
 
         trace = sampler(
-            chains=1,
+            chains=chains,
             random_seed=1322,
             keep_untransformed=True,
             postprocessing_backend=postprocessing_backend,
@@ -56,7 +66,7 @@ def test_transform_samples(sampler, postprocessing_backend):
     obs_at.set_value(-obs)
     with model:
         trace = sampler(
-            chains=2,
+            chains=chains,
             random_seed=1322,
             keep_untransformed=False,
             postprocessing_backend=postprocessing_backend,
@@ -73,6 +83,7 @@ def test_transform_samples(sampler, postprocessing_backend):
         sample_numpyro_nuts,
     ],
 )
+@pytest.mark.skipif(len(jax.devices()) < 2, reason="not enough devices")
 def test_deterministic_samples(sampler):
     aesara.config.on_opt_error = "raise"
     np.random.seed(13244)
@@ -207,7 +218,15 @@ def test_get_batched_jittered_initial_points():
     ],
 )
 @pytest.mark.parametrize("random_seed", (None, 123))
-@pytest.mark.parametrize("chains", (1, 2))
+@pytest.mark.parametrize(
+    "chains",
+    [
+        pytest.param(1),
+        pytest.param(
+            2, marks=pytest.mark.skipif(len(jax.devices()) < 2, reason="not enough devices")
+        ),
+    ],
+)
 def test_seeding(chains, random_seed, sampler):
     sample_kwargs = dict(
         tune=100,