Infer NumPyro InferenceData dims automatically (#2441)

* from_numpyro infers dims by default

* linted and formatted

* removed unneeded print statement

* caught edge case for inferring numpyro dims, changed failing test to expected behavior

* improved documentation

* better type hints

* updated changelog with az.from_numpyro change

* coords no longer required for infer_dims in NumPyroConverter. fixed infer_pred_dims logic

---------

Co-authored-by: kylejcaron <[email protected]>

Author: kylejcaron

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Unreleased
 
 ### New features
+- Make `arviz.from_numpyro(..., dims=None)` automatically infer dims from the numpyro model based on its numpyro.plate structure
 
 ### Maintenance and fixes
 

arviz/data/io_numpyro.py

@@ -1,7 +1,8 @@
 """NumPyro-specific conversion code."""
 
+from collections import defaultdict
 import logging
-from typing import Callable, Optional
+from typing import Any, Callable, Optional, Dict, List, Tuple
 
 import numpy as np
 
@@ -13,6 +14,70 @@
 _log = logging.getLogger(__name__)
 
 
+def _add_dims(dims_a: Dict[str, List[str]], dims_b: Dict[str, List[str]]) -> Dict[str, List[str]]:
+    merged = defaultdict(list)
+
+    for k, v in dims_a.items():
+        merged[k].extend(v)
+
+    for k, v in dims_b.items():
+        merged[k].extend(v)
+
+    # Convert back to a regular dict
+    return dict(merged)
+
+
+def infer_dims(
+    model: Callable,
+    model_args: Optional[Tuple[Any, ...]] = None,
+    model_kwargs: Optional[Dict[str, Any]] = None,
+) -> Dict[str, List[str]]:
+
+    from numpyro import handlers, distributions as dist
+    from numpyro.ops.pytree import PytreeTrace
+    from numpyro.infer.initialization import init_to_sample
+    import jax
+
+    model_args = tuple() if model_args is None else model_args
+    model_kwargs = dict() if model_args is None else model_kwargs
+
+    def _get_dist_name(fn):
+        if isinstance(fn, (dist.Independent, dist.ExpandedDistribution, dist.MaskedDistribution)):
+            return _get_dist_name(fn.base_dist)
+        return type(fn).__name__
+
+    def get_trace():
+        # We use `init_to_sample` to get around ImproperUniform distribution,
+        # which does not have `sample` method.
+        subs_model = handlers.substitute(
+            handlers.seed(model, 0),
+            substitute_fn=init_to_sample,
+        )
+        trace = handlers.trace(subs_model).get_trace(*model_args, **model_kwargs)
+        # Work around an issue where jax.eval_shape does not work
+        # for distribution output (e.g. the function `lambda: dist.Normal(0, 1)`)
+        # Here we will remove `fn` and store its name in the trace.
+        for _, site in trace.items():
+            if site["type"] == "sample":
+                site["fn_name"] = _get_dist_name(site.pop("fn"))
+            elif site["type"] == "deterministic":
+                site["fn_name"] = "Deterministic"
+        return PytreeTrace(trace)
+
+    # We use eval_shape to avoid any array computation.
+    trace = jax.eval_shape(get_trace).trace
+
+    named_dims = {}
+
+    for name, site in trace.items():
+        batch_dims = [frame.name for frame in sorted(site["cond_indep_stack"], key=lambda x: x.dim)]
+        event_dims = list(site.get("infer", {}).get("event_dims", []))
+        if site["type"] in ["sample", "deterministic"] and (batch_dims or event_dims):
+            named_dims[name] = batch_dims + event_dims
+
+    return named_dims
+
+
 class NumPyroConverter:
     """Encapsulate NumPyro specific logic."""
 
@@ -36,6 +101,7 @@ def __init__(
         coords=None,
         dims=None,
         pred_dims=None,
+        extra_event_dims=None,
         num_chains=1,
     ):
         """Convert NumPyro data into an InferenceData object.
@@ -58,9 +124,12 @@ def __init__(
         coords : dict[str] -> list[str]
             Map of dimensions to coordinates
         dims : dict[str] -> list[str]
-            Map variable names to their coordinates
+            Map variable names to their coordinates. Will be inferred if they are not provided.
         pred_dims: dict
             Dims for predictions data. Map variable names to their coordinates.
+        extra_event_dims: dict
+            Extra event dims for deterministic sites. Maps event dims that couldnt be inferred to
+            their coordinates.
         num_chains: int
             Number of chains used for sampling. Ignored if posterior is present.
         """
@@ -80,6 +149,7 @@ def __init__(
         self.coords = coords
         self.dims = dims
         self.pred_dims = pred_dims
+        self.extra_event_dims = extra_event_dims
         self.numpyro = numpyro
 
         def arbitrary_element(dct):
@@ -107,6 +177,10 @@ def arbitrary_element(dct):
             # model arguments and keyword arguments
             self._args = self.posterior._args  # pylint: disable=protected-access
             self._kwargs = self.posterior._kwargs  # pylint: disable=protected-access
+            self.dims = self.dims if self.dims is not None else self.infer_dims()
+            self.pred_dims = (
+                self.pred_dims if self.pred_dims is not None else self.infer_pred_dims()
+            )
         else:
             self.nchains = num_chains
             get_from = None
@@ -325,6 +399,23 @@ def to_inference_data(self):
             }
         )
 
+    @requires("posterior")
+    @requires("model")
+    def infer_dims(self) -> Dict[str, List[str]]:
+        dims = infer_dims(self.model, self._args, self._kwargs)
+        if self.extra_event_dims:
+            dims = _add_dims(dims, self.extra_event_dims)
+        return dims
+
+    @requires("posterior")
+    @requires("model")
+    @requires("predictions")
+    def infer_pred_dims(self) -> Dict[str, List[str]]:
+        dims = infer_dims(self.model, self._args, self._kwargs)
+        if self.extra_event_dims:
+            dims = _add_dims(dims, self.extra_event_dims)
+        return dims
+
 
 def from_numpyro(
     posterior=None,
@@ -339,10 +430,25 @@ def from_numpyro(
     coords=None,
     dims=None,
     pred_dims=None,
+    extra_event_dims=None,
     num_chains=1,
 ):
     """Convert NumPyro data into an InferenceData object.
 
+    If no dims are provided, this will infer batch dim names from NumPyro model plates.
+    For event dim names, such as with the ZeroSumNormal, `infer={"event_dims":dim_names}`
+    can be provided in numpyro.sample, i.e.::
+
+        # equivalent to dims entry, {"gamma": ["groups"]}
+        gamma = numpyro.sample(
+            "gamma",
+            dist.ZeroSumNormal(1, event_shape=(n_groups,)),
+            infer={"event_dims":["groups"]}
+        )
+
+    There is also an additional `extra_event_dims` input to cover any edge cases, for instance
+    deterministic sites with event dims (which dont have an `infer` argument to provide metadata).
+
     For a usage example read the
     :ref:`Creating InferenceData section on from_numpyro <creating_InferenceData>`
 
@@ -364,9 +470,10 @@ def from_numpyro(
     coords : dict[str] -> list[str]
         Map of dimensions to coordinates
     dims : dict[str] -> list[str]
-        Map variable names to their coordinates
+        Map variable names to their coordinates. Will be inferred if they are not provided.
     pred_dims: dict
-        Dims for predictions data. Map variable names to their coordinates.
+        Dims for predictions data. Map variable names to their coordinates. Default behavior is to
+        infer dims if this is not provided
     num_chains: int
         Number of chains used for sampling. Ignored if posterior is present.
     """
@@ -382,5 +489,6 @@ def from_numpyro(
         coords=coords,
         dims=dims,
         pred_dims=pred_dims,
+        extra_event_dims=extra_event_dims,
         num_chains=num_chains,
     ).to_inference_data()

arviz/tests/external_tests/test_data_numpyro.py

@@ -46,7 +46,9 @@ def predictions_data(self, data, predictions_params):
         )
         return predictions
 
-    def get_inference_data(self, data, eight_schools_params, predictions_data, predictions_params):
+    def get_inference_data(
+        self, data, eight_schools_params, predictions_data, predictions_params, infer_dims=False
+    ):
         posterior_samples = data.obj.get_samples()
         model = data.obj.sampler.model
         posterior_predictive = Predictive(model, posterior_samples)(
@@ -55,6 +57,12 @@ def get_inference_data(self, data, eight_schools_params, predictions_data, predi
         prior = Predictive(model, num_samples=500)(
             PRNGKey(2), eight_schools_params["J"], eight_schools_params["sigma"]
         )
+        dims = {"theta": ["school"], "eta": ["school"], "obs": ["school"]}
+        pred_dims = {"theta": ["school_pred"], "eta": ["school_pred"], "obs": ["school_pred"]}
+        if infer_dims:
+            dims = None
+            pred_dims = None
+
         predictions = predictions_data
         return from_numpyro(
             posterior=data.obj,
@@ -65,8 +73,8 @@ def get_inference_data(self, data, eight_schools_params, predictions_data, predi
                 "school": np.arange(eight_schools_params["J"]),
                 "school_pred": np.arange(predictions_params["J"]),
             },
-            dims={"theta": ["school"], "eta": ["school"], "obs": ["school"]},
-            pred_dims={"theta": ["school_pred"], "eta": ["school_pred"], "obs": ["school_pred"]},
+            dims=dims,
+            pred_dims=pred_dims,
         )
 
     def test_inference_data_namedtuple(self, data):
@@ -77,6 +85,7 @@ def test_inference_data_namedtuple(self, data):
         data.obj.get_samples = lambda *args, **kwargs: data_namedtuple
         inference_data = from_numpyro(
             posterior=data.obj,
+            dims={},  # This mock test needs to turn off autodims like so or mock group_by_chain
         )
         assert isinstance(data.obj.get_samples(), Samples)
         data.obj.get_samples = _old_fn
@@ -282,3 +291,121 @@ def model():
         mcmc.run(PRNGKey(0))
         inference_data = from_numpyro(mcmc)
         assert inference_data.observed_data
+
+    def test_mcmc_infer_dims(self):
+        import numpyro
+        import numpyro.distributions as dist
+        from numpyro.infer import MCMC, NUTS
+
+        def model():
+            # note: group2 gets assigned dim=-1 and group1 is assigned dim=-2
+            with numpyro.plate("group2", 5), numpyro.plate("group1", 10):
+                _ = numpyro.sample("param", dist.Normal(0, 1))
+
+        mcmc = MCMC(NUTS(model), num_warmup=10, num_samples=10)
+        mcmc.run(PRNGKey(0))
+        inference_data = from_numpyro(
+            mcmc, coords={"group1": np.arange(10), "group2": np.arange(5)}
+        )
+        assert inference_data.posterior.param.dims == ("chain", "draw", "group1", "group2")
+        assert all(dim in inference_data.posterior.param.coords for dim in ("group1", "group2"))
+
+    def test_mcmc_infer_unsorted_dims(self):
+        import numpyro
+        import numpyro.distributions as dist
+        from numpyro.infer import MCMC, NUTS
+
+        def model():
+            group1_plate = numpyro.plate("group1", 10, dim=-1)
+            group2_plate = numpyro.plate("group2", 5, dim=-2)
+
+            # the plate contexts are entered in a different order than the pre-defined dims
+            # we should make sure this still works because the trace has all of the info it needs
+            with group2_plate, group1_plate:
+                _ = numpyro.sample("param", dist.Normal(0, 1))
+
+        mcmc = MCMC(NUTS(model), num_warmup=10, num_samples=10)
+        mcmc.run(PRNGKey(0))
+        inference_data = from_numpyro(
+            mcmc, coords={"group1": np.arange(10), "group2": np.arange(5)}
+        )
+        assert inference_data.posterior.param.dims == ("chain", "draw", "group2", "group1")
+        assert all(dim in inference_data.posterior.param.coords for dim in ("group1", "group2"))
+
+    def test_mcmc_infer_dims_no_coords(self):
+        import numpyro
+        import numpyro.distributions as dist
+        from numpyro.infer import MCMC, NUTS
+
+        def model():
+            with numpyro.plate("group", 5):
+                _ = numpyro.sample("param", dist.Normal(0, 1))
+
+        mcmc = MCMC(NUTS(model), num_warmup=10, num_samples=10)
+        mcmc.run(PRNGKey(0))
+        inference_data = from_numpyro(mcmc)
+        assert inference_data.posterior.param.dims == ("chain", "draw", "group")
+
+    def test_mcmc_event_dims(self):
+        import numpyro
+        import numpyro.distributions as dist
+        from numpyro.infer import MCMC, NUTS
+
+        def model():
+            _ = numpyro.sample(
+                "gamma", dist.ZeroSumNormal(1, event_shape=(10,)), infer={"event_dims": ["groups"]}
+            )
+
+        mcmc = MCMC(NUTS(model), num_warmup=10, num_samples=10)
+        mcmc.run(PRNGKey(0))
+        inference_data = from_numpyro(mcmc, coords={"groups": np.arange(10)})
+        assert inference_data.posterior.gamma.dims == ("chain", "draw", "groups")
+        assert "groups" in inference_data.posterior.gamma.coords
+
+    @pytest.mark.xfail
+    def test_mcmc_inferred_dims_univariate(self):
+        import numpyro
+        import numpyro.distributions as dist
+        from numpyro.infer import MCMC, NUTS
+        import jax.numpy as jnp
+
+        def model():
+            alpha = numpyro.sample("alpha", dist.Normal(0, 1))
+            sigma = numpyro.sample("sigma", dist.HalfNormal(1))
+            with numpyro.plate("obs_idx", 3):
+                # mu is plated by obs_idx, but isnt broadcasted to the plate shape
+                # the expected behavior is that this should cause a failure
+                mu = numpyro.deterministic("mu", alpha)
+                return numpyro.sample("y", dist.Normal(mu, sigma), obs=jnp.array([-1, 0, 1]))
+
+        mcmc = MCMC(NUTS(model), num_warmup=10, num_samples=10)
+        mcmc.run(PRNGKey(0))
+        inference_data = from_numpyro(mcmc, coords={"obs_idx": np.arange(3)})
+        assert inference_data.posterior.mu.dims == ("chain", "draw", "obs_idx")
+        assert "obs_idx" in inference_data.posterior.mu.coords
+
+    def test_mcmc_extra_event_dims(self):
+        import numpyro
+        import numpyro.distributions as dist
+        from numpyro.infer import MCMC, NUTS
+
+        def model():
+            gamma = numpyro.sample("gamma", dist.ZeroSumNormal(1, event_shape=(10,)))
+            _ = numpyro.deterministic("gamma_plus1", gamma + 1)
+
+        mcmc = MCMC(NUTS(model), num_warmup=10, num_samples=10)
+        mcmc.run(PRNGKey(0))
+        inference_data = from_numpyro(
+            mcmc, coords={"groups": np.arange(10)}, extra_event_dims={"gamma_plus1": ["groups"]}
+        )
+        assert inference_data.posterior.gamma_plus1.dims == ("chain", "draw", "groups")
+        assert "groups" in inference_data.posterior.gamma_plus1.coords
+
+    def test_mcmc_predictions_infer_dims(
+        self, data, eight_schools_params, predictions_data, predictions_params
+    ):
+        inference_data = self.get_inference_data(
+            data, eight_schools_params, predictions_data, predictions_params, infer_dims=True
+        )
+        assert inference_data.predictions.obs.dims == ("chain", "draw", "J")
+        assert "J" in inference_data.predictions.obs.coords