Refactor sampling_jax

ricardoV94 · twiecki · commit 42482fdbb0ed · 2021-10-21T18:32:54.000+02:00
sampling_jax is no longer wrapped in a symbolic Op, making sampling more performant.
SharedVariables are now replaced by constants, and a ValueError is raised if any contains a default_update.
diff --git a/pymc/sampling_jax.py b/pymc/sampling_jax.py
@@ -15,144 +15,38 @@
 import pandas as pd
 
 from aesara.compile import SharedVariable
-from aesara.graph.basic import Apply, Constant, clone, graph_inputs
+from aesara.graph.basic import clone_replace, graph_inputs
 from aesara.graph.fg import FunctionGraph
-from aesara.graph.op import Op
 from aesara.graph.opt import MergeOptimizer
 from aesara.link.jax.dispatch import jax_funcify
-from aesara.tensor.type import TensorType
 
 from pymc import modelcontext
 from pymc.aesaraf import compile_rv_inplace
 
 warnings.warn("This module is experimental.")
 
 
-class NumPyroNUTS(Op):
-    def __init__(
-        self,
-        inputs,
-        outputs,
-        target_accept=0.8,
-        draws=1000,
-        tune=1000,
-        chains=4,
-        seed=None,
-        progress_bar=True,
-    ):
-        self.draws = draws
-        self.tune = tune
-        self.chains = chains
-        self.target_accept = target_accept
-        self.progress_bar = progress_bar
-        self.seed = seed
+def replace_shared_variables(graph):
+    """Replace shared variables in graph by their constant values
 
-        self.inputs, self.outputs = clone(inputs, outputs, copy_inputs=False)
-        self.inputs_type = tuple(input.type for input in inputs)
-        self.outputs_type = tuple(output.type for output in outputs)
-        self.nin = len(inputs)
-        self.nout = len(outputs)
-        self.nshared = len([v for v in inputs if isinstance(v, SharedVariable)])
-        self.samples_bcast = [self.chains == 1, self.draws == 1]
+    Raises
+    ------
+    ValueError
+        If any shared variable contains default_updates
+    """
 
-        self.fgraph = FunctionGraph(self.inputs, self.outputs, clone=False)
-        MergeOptimizer().optimize(self.fgraph)
+    shared_variables = [var for var in graph_inputs(graph) if isinstance(var, SharedVariable)]
 
-        super().__init__()
-
-    def make_node(self, *inputs):
-
-        # The samples for each variable
-        outputs = [
-            TensorType(v.dtype, self.samples_bcast + list(v.broadcastable))() for v in inputs
-        ]
-
-        # The leapfrog statistics
-        outputs += [TensorType("int64", self.samples_bcast)()]
-
-        all_inputs = list(inputs)
-        if self.nshared > 0:
-            all_inputs += self.inputs[-self.nshared :]
-
-        return Apply(self, all_inputs, outputs)
-
-    def do_constant_folding(self, *args):
-        return False
-
-    def perform(self, node, inputs, outputs):
-        raise NotImplementedError()
-
-
-@jax_funcify.register(NumPyroNUTS)
-def jax_funcify_NumPyroNUTS(op, node, **kwargs):
-    from numpyro.infer import MCMC, NUTS
-
-    draws = op.draws
-    tune = op.tune
-    chains = op.chains
-    target_accept = op.target_accept
-    progress_bar = op.progress_bar
-    seed = op.seed
-
-    # Compile the "inner" log-likelihood function.  This will have extra shared
-    # variable inputs as the last arguments
-    logp_fn = jax_funcify(op.fgraph, **kwargs)
-
-    if isinstance(logp_fn, (list, tuple)):
-        # This handles the new JAX backend, which always returns a tuple
-        logp_fn = logp_fn[0]
-
-    def _sample(*inputs):
-
-        if op.nshared > 0:
-            current_state = inputs[: -op.nshared]
-            shared_inputs = tuple(op.fgraph.inputs[-op.nshared :])
-        else:
-            current_state = inputs
-            shared_inputs = ()
-
-        def log_fn_wrap(x):
-            res = logp_fn(
-                *(
-                    x
-                    # We manually obtain the shared values and added them
-                    # as arguments to our compiled "inner" function
-                    + tuple(
-                        v.get_value(borrow=True, return_internal_type=True) for v in shared_inputs
-                    )
-                )
-            )
-
-            if isinstance(res, (list, tuple)):
-                # This handles the new JAX backend, which always returns a tuple
-                res = res[0]
-
-            return -res
-
-        nuts_kernel = NUTS(
-            potential_fn=log_fn_wrap,
-            target_accept_prob=target_accept,
-            adapt_step_size=True,
-            adapt_mass_matrix=True,
-            dense_mass=False,
-        )
-
-        pmap_numpyro = MCMC(
-            nuts_kernel,
-            num_warmup=tune,
-            num_samples=draws,
-            num_chains=chains,
-            postprocess_fn=None,
-            chain_method="parallel",
-            progress_bar=progress_bar,
+    if any(hasattr(var, "default_update") for var in shared_variables):
+        raise ValueError(
+            "Graph contains shared variables with default_update which cannot "
+            "be safely replaced."
         )
 
-        pmap_numpyro.run(seed, init_params=current_state, extra_fields=("num_steps",))
-        samples = pmap_numpyro.get_samples(group_by_chain=True)
-        leapfrogs_taken = pmap_numpyro.get_extra_fields(group_by_chain=True)["num_steps"]
-        return tuple(samples) + (leapfrogs_taken,)
+    replacements = {var: at.constant(var.get_value(borrow=True)) for var in shared_variables}
 
-    return _sample
+    new_graph = clone_replace(graph, replace=replacements)
+    return new_graph
 
 
 def sample_numpyro_nuts(
@@ -165,72 +59,101 @@ def sample_numpyro_nuts(
     progress_bar=True,
     keep_untransformed=False,
 ):
+    from numpyro.infer import MCMC, NUTS
+
     model = modelcontext(model)
 
-    seed = jax.random.PRNGKey(random_seed)
+    tic1 = pd.Timestamp.now()
+    print("Compiling...", file=sys.stdout)
 
     rv_names = [rv.name for rv in model.value_vars]
     init_state = [model.initial_point[rv_name] for rv_name in rv_names]
     init_state_batched = jax.tree_map(lambda x: np.repeat(x[None, ...], chains, axis=0), init_state)
-    init_state_batched_at = [at.as_tensor(v) for v in init_state_batched]
 
-    nuts_inputs = sorted(
-        (v for v in graph_inputs([model.logpt]) if not isinstance(v, Constant)),
-        key=lambda x: isinstance(x, SharedVariable),
-    )
-    map_seed = jax.random.split(seed, chains)
-    numpyro_samples = NumPyroNUTS(
-        nuts_inputs,
-        [model.logpt],
-        target_accept=target_accept,
-        draws=draws,
-        tune=tune,
-        chains=chains,
-        seed=map_seed,
-        progress_bar=progress_bar,
-    )(*init_state_batched_at)
+    logpt = replace_shared_variables([model.logpt])[0]
+    logpt_fgraph = FunctionGraph(outputs=[logpt], clone=False)
+    MergeOptimizer().optimize(logpt_fgraph)
+    logp_fn = jax_funcify(logpt_fgraph)
 
-    # Un-transform the transformed variables in JAX
-    sample_outputs = []
-    for i, (value_var, rv_samples) in enumerate(zip(model.value_vars, numpyro_samples[:-1])):
-        rv = model.values_to_rvs[value_var]
-        transform = getattr(value_var.tag, "transform", None)
-        if transform is not None:
-            untrans_value_var = transform.backward(rv, rv_samples)
-            untrans_value_var.name = rv.name
-            sample_outputs.append(untrans_value_var)
+    if isinstance(logp_fn, (list, tuple)):
+        # This handles the new JAX backend, which always returns a tuple
+        logp_fn = logp_fn[0]
 
-            if keep_untransformed:
-                rv_samples.name = value_var.name
-                sample_outputs.append(rv_samples)
-        else:
-            rv_samples.name = rv.name
-            sample_outputs.append(rv_samples)
+    def logp_fn_wrap(x):
+        res = logp_fn(*x)
 
-    print("Compiling...", file=sys.stdout)
+        if isinstance(res, (list, tuple)):
+            # This handles the new JAX backend, which always returns a tuple
+            res = res[0]
 
-    tic1 = pd.Timestamp.now()
-    _sample = compile_rv_inplace(
-        [],
-        sample_outputs + [numpyro_samples[-1]],
-        allow_input_downcast=True,
-        on_unused_input="ignore",
-        accept_inplace=True,
-        mode="JAX",
+        # Jax expects a potential with the opposite sign of model.logpt
+        return -res
+
+    nuts_kernel = NUTS(
+        potential_fn=logp_fn_wrap,
+        target_accept_prob=target_accept,
+        adapt_step_size=True,
+        adapt_mass_matrix=True,
+        dense_mass=False,
+    )
+
+    pmap_numpyro = MCMC(
+        nuts_kernel,
+        num_warmup=tune,
+        num_samples=draws,
+        num_chains=chains,
+        postprocess_fn=None,
+        chain_method="parallel",
+        progress_bar=progress_bar,
     )
-    tic2 = pd.Timestamp.now()
 
+    tic2 = pd.Timestamp.now()
     print("Compilation time = ", tic2 - tic1, file=sys.stdout)
 
     print("Sampling...", file=sys.stdout)
 
-    *mcmc_samples, leapfrogs_taken = _sample()
-    tic3 = pd.Timestamp.now()
+    seed = jax.random.PRNGKey(random_seed)
+    map_seed = jax.random.split(seed, chains)
 
+    pmap_numpyro.run(map_seed, init_params=init_state_batched, extra_fields=("num_steps",))
+    raw_mcmc_samples = pmap_numpyro.get_samples(group_by_chain=True)
+
+    tic3 = pd.Timestamp.now()
     print("Sampling time = ", tic3 - tic2, file=sys.stdout)
 
-    posterior = {k.name: v for k, v in zip(sample_outputs, mcmc_samples)}
+    print("Transforming variables...", file=sys.stdout)
+    mcmc_samples = []
+    for i, (value_var, raw_samples) in enumerate(zip(model.value_vars, raw_mcmc_samples)):
+        raw_samples = at.constant(np.asarray(raw_samples))
+
+        rv = model.values_to_rvs[value_var]
+        transform = getattr(value_var.tag, "transform", None)
+
+        if transform is not None:
+            # TODO: This will fail when the transformation depends on another variable
+            #  such as in interval transform with RVs as edges
+            trans_samples = transform.backward(rv, raw_samples)
+            trans_samples.name = rv.name
+            mcmc_samples.append(trans_samples)
+
+            if keep_untransformed:
+                raw_samples.name = value_var.name
+                mcmc_samples.append(raw_samples)
+        else:
+            raw_samples.name = rv.name
+            mcmc_samples.append(raw_samples)
+
+    mcmc_varnames = [var.name for var in mcmc_samples]
+    mcmc_samples = compile_rv_inplace(
+        [],
+        mcmc_samples,
+        mode="JAX",
+    )()
+
+    tic4 = pd.Timestamp.now()
+    print("Transformation time = ", tic4 - tic3, file=sys.stdout)
 
+    posterior = {k: v for k, v in zip(mcmc_varnames, mcmc_samples)}
     az_trace = az.from_dict(posterior=posterior)
 
     return az_trace
diff --git a/pymc/tests/test_sampling_jax.py b/pymc/tests/test_sampling_jax.py
@@ -1,9 +1,13 @@
 import aesara
 import numpy as np
+import pytest
+
+from aesara.compile import SharedVariable
+from aesara.graph import graph_inputs
 
 import pymc as pm
 
-from pymc.sampling_jax import sample_numpyro_nuts
+from pymc.sampling_jax import replace_shared_variables, sample_numpyro_nuts
 
 
 def test_transform_samples():
@@ -29,7 +33,20 @@ def test_transform_samples():
 
     obs_at.set_value(-obs)
     with model:
-        trace = sample_numpyro_nuts(chains=1, random_seed=1322, keep_untransformed=False)
+        trace = sample_numpyro_nuts(chains=2, random_seed=1322, keep_untransformed=False)
 
     assert -11 < trace.posterior["a"].mean() < -8
     assert 1.5 < trace.posterior["sigma"].mean() < 2.5
+
+
+def test_replace_shared_variables():
+
+    x = aesara.shared(5, name="shared_x")
+
+    new_x = replace_shared_variables([x])
+    shared_variables = [var for var in graph_inputs(new_x) if isinstance(var, SharedVariable)]
+    assert not shared_variables
+
+    x.default_update = x + 1
+    with pytest.raises(ValueError, match="shared variables with default_update"):
+        replace_shared_variables([x])
