Repository: arviz-devs/arviz

Issue:
Expose skip_dims/flatten argument
All density plots (except plot_ppc) reduce the chain and draw dimensions, some have a combined argument to avoid reducing the chain dimension,  but in some cases this is not enough, and similarly to plot_ppc behaviour, we would like to choose at will which dimensions to reduce. I think we should support that. 

## Thoughts on implementation
Exposing skip_dims argument from xarray_var_iter should do the trick. This is basically what plot_ppc does with its flatten argument


Discussion:
OriolAbril: As a way of showing how this feature will have immediate use and application, [here](https://discourse.mc-stan.org/t/ppcs-in-arviz/22176) is one question where we could use this to mimic `ppc_stat_2d` in bayesplot with our `pair_plot`, so thanks again for the PR! :rocket: 
codecov[bot]: # [Codecov](https://codecov.io/gh/arviz-devs/arviz/pull/1676?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=arviz-devs) Report
> Merging [#1676](https://codecov.io/gh/arviz-devs/arviz/pull/1676?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=arviz-devs) (e25bc59) into [main](https://codecov.io/gh/arviz-devs/arviz/commit/31c0cdac56318afa4c4df5e65d16bdb90065baf7?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=arviz-devs) (31c0cda) will **increase** coverage by `0.01%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/arviz-devs/arviz/pull/1676/graphs/tree.svg?width=650&height=150&src=pr&token=KZTdQCEf7t&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=arviz-devs)](https://codecov.io/gh/arviz-devs/arviz/pull/1676?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=arviz-devs)

```diff
@@            Coverage Diff             @@
##             main    #1676      +/-   ##
==========================================
+ Coverage   90.83%   90.84%   +0.01%     
==========================================
  Files         114      114              
  Lines       12311    12327      +16     
==========================================
+ Hits        11183    11199      +16     
  Misses       1128     1128              
```


| [Impacted Files](https://codecov.io/gh/arviz-devs/arviz/pull/1676?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=arviz-devs) | Coverage Δ | |
|---|---|---|
| [arviz/plots/distcomparisonplot.py](https://codecov.io/gh/arviz-devs/arviz/pull/1676/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=arviz-devs#diff-YXJ2aXovcGxvdHMvZGlzdGNvbXBhcmlzb25wbG90LnB5) | `93.02% <ø> (ø)` | |
| [arviz/plots/pairplot.py](https://codecov.io/gh/arviz-devs/arviz/pull/1676/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=arviz-devs#diff-YXJ2aXovcGxvdHMvcGFpcnBsb3QucHk=) | `88.88% <ø> (ø)` | |
| [arviz/plots/posteriorplot.py](https://codecov.io/gh/arviz-devs/arviz/pull/1676/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=arviz-devs#diff-YXJ2aXovcGxvdHMvcG9zdGVyaW9ycGxvdC5weQ==) | `91.66% <ø> (ø)` | |
| [arviz/plots/rankplot.py](https://codecov.io/gh/arviz-devs/arviz/pull/1676/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=arviz-devs#diff-YXJ2aXovcGxvdHMvcmFua3Bsb3QucHk=) | `91.42% <ø> (ø)` | |
| [arviz/plots/violinplot.py](https://codecov.io/gh/arviz-devs/arviz/pull/1676/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=arviz-devs#diff-YXJ2aXovcGxvdHMvdmlvbGlucGxvdC5weQ==) | `88.46% <ø> (ø)` | |
| [arviz/plots/backends/bokeh/forestplot.py](https://codecov.io/gh/arviz-devs/arviz/pull/1676/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=arviz-devs#diff-YXJ2aXovcGxvdHMvYmFja2VuZHMvYm9rZWgvZm9yZXN0cGxvdC5weQ==) | `94.94% <100.00%> (+0.02%)` | :arrow_up: |
| [arviz/plots/backends/matplotlib/forestplot.py](https://codecov.io/gh/arviz-devs/arviz/pull/1676/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=arviz-devs#diff-YXJ2aXovcGxvdHMvYmFja2VuZHMvbWF0cGxvdGxpYi9mb3Jlc3RwbG90LnB5) | `97.86% <100.00%> (+0.01%)` | :arrow_up: |
| [arviz/plots/densityplot.py](https://codecov.io/gh/arviz-devs/arviz/pull/1676/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=arviz-devs#diff-YXJ2aXovcGxvdHMvZGVuc2l0eXBsb3QucHk=) | `91.07% <100.00%> (ø)` | |
| [arviz/plots/forestplot.py](https://codecov.io/gh/arviz-devs/arviz/pull/1676/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=arviz-devs#diff-YXJ2aXovcGxvdHMvZm9yZXN0cGxvdC5weQ==) | `94.87% <100.00%> (+0.27%)` | :arrow_up: |
| [arviz/plots/parallelplot.py](https://codecov.io/gh/arviz-devs/arviz/pull/1676/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=arviz-devs#diff-YXJ2aXovcGxvdHMvcGFyYWxsZWxwbG90LnB5) | `97.77% <100.00%> (ø)` | |
| ... and [1 more](https://codecov.io/gh/arviz-devs/arviz/pull/1676/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=arviz-devs) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/arviz-devs/arviz/pull/1676?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=arviz-devs).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=arviz-devs)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/arviz-devs/arviz/pull/1676?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=arviz-devs). Last update [31c0cda...e25bc59](https://codecov.io/gh/arviz-devs/arviz/pull/1676?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=arviz-devs). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=arviz-devs).


Patch:
diff --git a/arviz/plots/backends/bokeh/forestplot.py b/arviz/plots/backends/bokeh/forestplot.py
index f8986bf..043201c 100644
--- a/arviz/plots/backends/bokeh/forestplot.py
+++ b/arviz/plots/backends/bokeh/forestplot.py
@@ -34,6 +34,7 @@ def plot_forest(
     var_names,
     model_names,
     combined,
+    combine_dims,
     colors,
     figsize,
     width_ratios,
@@ -64,6 +65,7 @@ def plot_forest(
         var_names=var_names,
         model_names=model_names,
         combined=combined,
+        combine_dims=combine_dims,
         colors=colors,
         labeller=labeller,
     )
@@ -214,7 +216,7 @@ class PlotHandler:
 
     # pylint: disable=inconsistent-return-statements
 
-    def __init__(self, datasets, var_names, model_names, combined, colors, labeller):
+    def __init__(self, datasets, var_names, model_names, combined, combine_dims, colors, labeller):
         self.data = datasets
 
         if model_names is None:
@@ -240,6 +242,7 @@ class PlotHandler:
             self.var_names = list(reversed(var_names))  # y-values are upside down
 
         self.combined = combined
+        self.combine_dims = combine_dims
 
         if colors == "cycle":
             colors = [
@@ -266,6 +269,7 @@ class PlotHandler:
                 y,
                 model_names=self.model_names,
                 combined=self.combined,
+                combine_dims=self.combine_dims,
                 colors=self.colors,
                 labeller=self.labeller,
             )
@@ -616,12 +620,15 @@ class PlotHandler:
 class VarHandler:
     """Handle individual variable logic."""
 
-    def __init__(self, var_name, data, y_start, model_names, combined, colors, labeller):
+    def __init__(
+        self, var_name, data, y_start, model_names, combined, combine_dims, colors, labeller
+    ):
         self.var_name = var_name
         self.data = data
         self.y_start = y_start
         self.model_names = model_names
         self.combined = combined
+        self.combine_dims = combine_dims
         self.colors = colors
         self.labeller = labeller
         self.model_color = dict(zip(self.model_names, self.colors))
@@ -634,10 +641,10 @@ class VarHandler:
         """Iterate over models and chains for each variable."""
         if self.combined:
             grouped_data = [[(0, datum)] for datum in self.data]
-            skip_dims = {"chain"}
+            skip_dims = self.combine_dims.union({"chain"})
         else:
             grouped_data = [datum.groupby("chain") for datum in self.data]
-            skip_dims = set()
+            skip_dims = self.combine_dims
 
         label_dict = OrderedDict()
         selection_list = []
diff --git a/arviz/plots/backends/matplotlib/forestplot.py b/arviz/plots/backends/matplotlib/forestplot.py
index d19aaaa..8384603 100644
--- a/arviz/plots/backends/matplotlib/forestplot.py
+++ b/arviz/plots/backends/matplotlib/forestplot.py
@@ -29,6 +29,7 @@ def plot_forest(
     var_names,
     model_names,
     combined,
+    combine_dims,
     colors,
     figsize,
     width_ratios,
@@ -59,6 +60,7 @@ def plot_forest(
         var_names=var_names,
         model_names=model_names,
         combined=combined,
+        combine_dims=combine_dims,
         colors=colors,
         labeller=labeller,
     )
@@ -175,7 +177,7 @@ class PlotHandler:
 
     # pylint: disable=inconsistent-return-statements
 
-    def __init__(self, datasets, var_names, model_names, combined, colors, labeller):
+    def __init__(self, datasets, var_names, model_names, combined, combine_dims, colors, labeller):
         self.data = datasets
 
         if model_names is None:
@@ -201,6 +203,7 @@ class PlotHandler:
             self.var_names = list(reversed(var_names))  # y-values are upside down
 
         self.combined = combined
+        self.combine_dims = combine_dims
 
         if colors == "cycle":
             # TODO: Use matplotlib prop cycle instead
@@ -223,6 +226,7 @@ class PlotHandler:
                 y,
                 model_names=self.model_names,
                 combined=self.combined,
+                combine_dims=self.combine_dims,
                 colors=self.colors,
                 labeller=self.labeller,
             )
@@ -510,12 +514,15 @@ class PlotHandler:
 class VarHandler:
     """Handle individual variable logic."""
 
-    def __init__(self, var_name, data, y_start, model_names, combined, colors, labeller):
+    def __init__(
+        self, var_name, data, y_start, model_names, combined, combine_dims, colors, labeller
+    ):
         self.var_name = var_name
         self.data = data
         self.y_start = y_start
         self.model_names = model_names
         self.combined = combined
+        self.combine_dims = combine_dims
         self.colors = colors
         self.labeller = labeller
         self.model_color = dict(zip(self.model_names, self.colors))
@@ -528,10 +535,10 @@ class VarHandler:
         """Iterate over models and chains for each variable."""
         if self.combined:
             grouped_data = [[(0, datum)] for datum in self.data]
-            skip_dims = {"chain"}
+            skip_dims = self.combine_dims.union({"chain"})
         else:
             grouped_data = [datum.groupby("chain") for datum in self.data]
-            skip_dims = set()
+            skip_dims = self.combine_dims
 
         label_dict = OrderedDict()
         selection_list = []
diff --git a/arviz/plots/densityplot.py b/arviz/plots/densityplot.py
index dd881a8..ee1593f 100644
--- a/arviz/plots/densityplot.py
+++ b/arviz/plots/densityplot.py
@@ -18,6 +18,7 @@ def plot_density(
     data_labels=None,
     var_names=None,
     filter_vars=None,
+    combine_dims=None,
     transform=None,
     hdi_prob=None,
     point_estimate="auto",
@@ -64,6 +65,9 @@ def plot_density(
         interpret var_names as substrings of the real variables names. If "regex",
         interpret var_names as regular expressions on the real variables names. A la
         ``pandas.filter``.
+    combine_dims : set_like of str, optional
+        List of dimensions to reduce. Defaults to reducing only the "chain" and "draw" dimensions.
+        See the :ref:`this section <common_combine_dims>` for usage examples.
     transform : callable
         Function to transform data (defaults to None i.e. the identity function)
     hdi_prob : float
@@ -193,6 +197,7 @@ def plot_density(
         labeller = BaseLabeller()
 
     var_names = _var_names(var_names, datasets, filter_vars)
+
     n_data = len(datasets)
 
     if data_labels is None:
@@ -212,7 +217,10 @@ def plot_density(
         if not 1 >= hdi_prob > 0:
             raise ValueError("The value of hdi_prob should be in the interval (0, 1]")
 
-    to_plot = [list(xarray_var_iter(data, var_names, combined=True)) for data in datasets]
+    to_plot = [
+        list(xarray_var_iter(data, var_names, combined=True, skip_dims=combine_dims))
+        for data in datasets
+    ]
     all_labels = []
     length_plotters = []
     for plotters in to_plot:
diff --git a/arviz/plots/distcomparisonplot.py b/arviz/plots/distcomparisonplot.py
index 4a7c63b..e3c115e 100644
--- a/arviz/plots/distcomparisonplot.py
+++ b/arviz/plots/distcomparisonplot.py
@@ -13,6 +13,7 @@ def plot_dist_comparison(
     textsize=None,
     var_names=None,
     coords=None,
+    combine_dims=None,
     transform=None,
     legend=True,
     labeller=None,
@@ -50,6 +51,10 @@ def plot_dist_comparison(
         Dictionary mapping dimensions to selected coordinates to be plotted.
         Dimensions without a mapping specified will include all coordinates for
         that dimension.
+    combine_dims : set_like of str, optional
+        List of dimensions to reduce. Defaults to reducing only the "chain" and "draw" dimensions.
+        See the :ref:`this section <common_combine_dims>` for usage examples.
+
     transform : callable
         Function to transform data (defaults to None i.e. the identity function)
     legend : bool
@@ -136,7 +141,9 @@ def plot_dist_comparison(
     len_plots = rcParams["plot.max_subplots"] // (len(groups) + 1)
     len_plots = len_plots if len_plots else 1
     dc_plotters = [
-        list(xarray_var_iter(data, var_names=var, combined=True))[:len_plots]
+        list(xarray_var_iter(data, var_names=var, combined=True, skip_dims=combine_dims))[
+            :len_plots
+        ]
         for data, var in zip(datasets, var_names)
     ]
 
diff --git a/arviz/plots/forestplot.py b/arviz/plots/forestplot.py
index d5e4224..c5017d2 100644
--- a/arviz/plots/forestplot.py
+++ b/arviz/plots/forestplot.py
@@ -15,6 +15,7 @@ def plot_forest(
     transform=None,
     coords=None,
     combined=False,
+    combine_dims=None,
     hdi_prob=None,
     rope=None,
     quartiles=True,
@@ -55,6 +56,9 @@ def plot_forest(
     var_names: list[str], optional
         List of variables to plot (defaults to None, which results in all variables plotted)
         Prefix the variables by ``~`` when you want to exclude them from the plot.
+    combine_dims : set_like of str, optional
+        List of dimensions to reduce. Defaults to reducing only the "chain" and "draw" dimensions.
+        See the :ref:`this section <common_combine_dims>` for usage examples.
     filter_vars: {None, "like", "regex"}, optional, default=None
         If None(default), interpret var_names as the real variables names. If "like", interpret
         var_names as substrings of the real variables names. If "regex", interpret var_names as
@@ -219,6 +223,9 @@ def plot_forest(
     if coords is None:
         coords = {}
 
+    if combine_dims is None:
+        combine_dims = set()
+
     if labeller is None:
         labeller = NoModelLabeller() if legend else BaseLabeller()
 
@@ -253,6 +260,7 @@ def plot_forest(
         var_names=var_names,
         model_names=model_names,
         combined=combined,
+        combine_dims=combine_dims,
         colors=colors,
         figsize=figsize,
         width_ratios=width_ratios,
diff --git a/arviz/plots/pairplot.py b/arviz/plots/pairplot.py
index d687c08..4fe0422 100644
--- a/arviz/plots/pairplot.py
+++ b/arviz/plots/pairplot.py
@@ -17,6 +17,7 @@ def plot_pair(
     group="posterior",
     var_names: Optional[List[str]] = None,
     filter_vars: Optional[str] = None,
+    combine_dims=None,
     coords=None,
     marginals=False,
     figsize=None,
@@ -62,6 +63,9 @@ def plot_pair(
         interpret var_names as substrings of the real variables names. If "regex",
         interpret var_names as regular expressions on the real variables names. A la
         ``pandas.filter``.
+    combine_dims : set_like of str, optional
+        List of dimensions to reduce. Defaults to reducing only the "chain" and "draw" dimensions.
+        See the :ref:`this section <common_combine_dims>` for usage examples.
     coords: mapping, optional
         Coordinates of var_names to be plotted. Passed to :meth:`xarray.Dataset.sel`.
     marginals: bool, optional
@@ -211,7 +215,9 @@ def plot_pair(
     dataset = convert_to_dataset(data, group=group)
     var_names = _var_names(var_names, dataset, filter_vars)
     plotters = list(
-        xarray_var_iter(get_coords(dataset, coords), var_names=var_names, combined=True)
+        xarray_var_iter(
+            get_coords(dataset, coords), var_names=var_names, skip_dims=combine_dims, combined=True
+        )
     )
     flat_var_names = [
         labeller.make_label_vert(var_name, sel, isel) for var_name, sel, isel, _ in plotters
diff --git a/arviz/plots/parallelplot.py b/arviz/plots/parallelplot.py
index 96a8941..04ec286 100644
--- a/arviz/plots/parallelplot.py
+++ b/arviz/plots/parallelplot.py
@@ -124,7 +124,11 @@ def plot_parallel(
 
     # Get diverging draws and combine chains
     divergent_data = convert_to_dataset(data, group="sample_stats")
-    _, diverging_mask = xarray_to_ndarray(divergent_data, var_names=("diverging",), combined=True)
+    _, diverging_mask = xarray_to_ndarray(
+        divergent_data,
+        var_names=("diverging",),
+        combined=True,
+    )
     diverging_mask = np.squeeze(diverging_mask)
 
     # Get posterior draws and combine chains
diff --git a/arviz/plots/posteriorplot.py b/arviz/plots/posteriorplot.py
index 516bd68..0e7e481 100644
--- a/arviz/plots/posteriorplot.py
+++ b/arviz/plots/posteriorplot.py
@@ -11,6 +11,7 @@ def plot_posterior(
     data,
     var_names=None,
     filter_vars=None,
+    combine_dims=None,
     transform=None,
     coords=None,
     grid=None,
@@ -52,6 +53,9 @@ def plot_posterior(
         interpret var_names as substrings of the real variables names. If "regex",
         interpret var_names as regular expressions on the real variables names. A la
         ``pandas.filter``.
+    combine_dims : set_like of str, optional
+        List of dimensions to reduce. Defaults to reducing only the "chain" and "draw" dimensions.
+        See the :ref:`this section <common_combine_dims>` for usage examples.
     transform: callable
         Function to transform data (defaults to None i.e.the identity function)
     coords: mapping, optional
@@ -246,7 +250,11 @@ def plot_posterior(
         kind = rcParams["plot.density_kind"]
 
     plotters = filter_plotters_list(
-        list(xarray_var_iter(get_coords(data, coords), var_names=var_names, combined=True)),
+        list(
+            xarray_var_iter(
+                get_coords(data, coords), var_names=var_names, combined=True, skip_dims=combine_dims
+            )
+        ),
         "plot_posterior",
     )
     length_plotters = len(plotters)
diff --git a/arviz/plots/rankplot.py b/arviz/plots/rankplot.py
index 7df8f61..0b5eb97 100644
--- a/arviz/plots/rankplot.py
+++ b/arviz/plots/rankplot.py
@@ -169,7 +169,14 @@ def plot_rank(
         posterior_data = posterior_data.sel(**coords)
     var_names = _var_names(var_names, posterior_data, filter_vars)
     plotters = filter_plotters_list(
-        list(xarray_var_iter(posterior_data, var_names=var_names, combined=True)), "plot_rank"
+        list(
+            xarray_var_iter(
+                posterior_data,
+                var_names=var_names,
+                combined=True,
+            )
+        ),
+        "plot_rank",
     )
     length_plotters = len(plotters)
 
diff --git a/arviz/plots/violinplot.py b/arviz/plots/violinplot.py
index e692763..37f2b4d 100644
--- a/arviz/plots/violinplot.py
+++ b/arviz/plots/violinplot.py
@@ -10,6 +10,7 @@ from .plot_utils import default_grid, filter_plotters_list, get_plotting_functio
 def plot_violin(
     data,
     var_names=None,
+    combine_dims=None,
     filter_vars=None,
     transform=None,
     quartiles=True,
@@ -45,6 +46,9 @@ def plot_violin(
     var_names: list of variable names, optional
         Variables to be plotted, if None all variable are plotted. Prefix the
         variables by ``~`` when you want to exclude them from the plot.
+    combine_dims : set_like of str, optional
+        List of dimensions to reduce. Defaults to reducing only the "chain" and "draw" dimensions.
+        See the :ref:`this section <common_combine_dims>` for usage examples.
     filter_vars: {None, "like", "regex"}, optional, default=None
         If `None` (default), interpret var_names as the real variables names. If "like",
         interpret var_names as substrings of the real variables names. If "regex",
@@ -135,7 +139,8 @@ def plot_violin(
     var_names = _var_names(var_names, data, filter_vars)
 
     plotters = filter_plotters_list(
-        list(xarray_var_iter(data, var_names=var_names, combined=True)), "plot_violin"
+        list(xarray_var_iter(data, var_names=var_names, combined=True, skip_dims=combine_dims)),
+        "plot_violin",
     )
 
     rows, cols = default_grid(len(plotters), grid=grid)
diff --git a/doc/source/user_guide/plots_arguments_guide.md b/doc/source/user_guide/plots_arguments_guide.md
index aa11bc8..adfac4c 100644
--- a/doc/source/user_guide/plots_arguments_guide.md
+++ b/doc/source/user_guide/plots_arguments_guide.md
@@ -17,8 +17,8 @@ Most of these plotting functions have common arguments. These common arguments a
 
 ```{code-cell}
 import arviz as az
-data = az.load_arviz_data('centered_eight');
-non_centered = az.load_arviz_data('non_centered_eight');
+centered_eight = az.load_arviz_data('centered_eight');
+non_centered_eight = az.load_arviz_data('non_centered_eight');
 ```
 
 :::{warning} Page in construction
@@ -32,24 +32,24 @@ Variables to be plotted, if None all variables are plotted. Prefix the variables
 Plot default autocorrelation
 
 ```{code-cell}
-az.plot_posterior(data);
+az.plot_posterior(centered_eight);
 ```
 
 Plot one variable by setting `var_names=var1`
 
 ```{code-cell}
-az.plot_posterior(data, var_names=['mu']);
+az.plot_posterior(centered_eight, var_names=['mu']);
 ```
 
 Plot subset variables by specifying variable name exactly
 
 ```{code-cell}
-az.plot_posterior(data, var_names=['mu', 'tau']);
+az.plot_posterior(centered_eight, var_names=['mu', 'tau']);
 ```
 
 Plot variables with regular expressions
 ```{code-cell}
-az.plot_posterior(data, var_names=['mu', '^the'], filter_vars="regex");
+az.plot_posterior(centered_eight, var_names=['mu', '^the'], filter_vars="regex");
 ```
 
 (common_filter_vars)=
@@ -59,7 +59,7 @@ If None (default), interpret `var_names` as the real variables names.
 Plot using the default value of `filter_vars` which is `None`
 
 ```{code-cell}
-az.plot_posterior(data, var_names=['mu']);
+az.plot_posterior(centered_eight, var_names=['mu']);
 ```
 
 If “like”, interpret `var_names` as substrings of the real variables names.
@@ -67,7 +67,7 @@ If “like”, interpret `var_names` as substrings of the real variables names.
 Plot using `filter_vars="like"`
 
 ```{code-cell}
-az.plot_posterior(data, var_names=['mu', '^the'], filter_vars="like");
+az.plot_posterior(centered_eight, var_names=['mu', 'the'], filter_vars="like");
 ```
 
 If “regex”, interpret `var_names` as regular expressions on the real variables names. A la `pandas.filter`.
@@ -75,7 +75,7 @@ If “regex”, interpret `var_names` as regular expressions on the real variabl
 Plot using `filter_vars="regex"`
 
 ```{code-cell}
-az.plot_posterior(data, var_names=['mu', '^the'], filter_vars="regex");
+az.plot_posterior(centered_eight, var_names=['mu', '^the'], filter_vars="regex");
 ```
 
 (common_coords)=
@@ -86,7 +86,44 @@ Using coords argument to plot only a subset of data
 
 ```{code-cell}
 coords = {"school": ["Choate","Phillips Exeter"]};
-az.plot_posterior(data, var_names=["mu", "theta"], coords=coords);
+az.plot_posterior(centered_eight, var_names=["mu", "theta"], coords=coords);
+```
+
+(common_combine_dims)=
+## `combine_dims`
+Set like argument containing dimensions to reduce for plots that represent probability
+distributions. When used, the dimensions on `combine_dims` are added to `chain` and
+`draw` dimensions and reduced to generate a single KDE/histogram/dot plot from >2D
+arrays.
+
+`combine_dims` dims can be used to generate the KDE of the distribution of `theta`
+when combining all schools together.
+
+```{code-cell}
+az.plot_posterior(centered_eight, var_names=["mu", "theta"], combine_dims={"school"})
+```
+
+:::{warning}
+`plot_pair` also supports `combine_dims` argument, but it's the users responsibility
+to ensure the variables to be plotted have compatible dimensions when reducing
+them pairwise.
+:::
+
+Both `theta` and `theta_t` have the `school` dimension, so we can use `combine_dims`
+in `plot_pair` to generate their global 2d KDE.
+
+```{code-cell}
+az.plot_pair(
+    non_centered_eight, var_names=["theta", "theta_t"], combine_dims={"school"}, kind="kde"
+)
+```
+
+`mu` however does not, so trying to plot their combined distribution errors out:
+
+```{code-cell}
+:tags: [raises-exception]
+
+az.plot_pair(non_centered_eight, var_names=["theta", "mu"], combine_dims={"school"})
 ```
 
 (common_hdi_prob)=
@@ -102,7 +139,7 @@ Number of rows and columns. Defaults to None, the rows and columns are automatic
 Plot variables in a 4x5 grid
 
 ```{code-cell}
-az.plot_density([data, non_centered], grid=(4, 5));
+az.plot_density([centered_eight, non_centered_eight], grid=(4, 5));
 ```
 
 (common_figsize)=
diff --git a/examples/bokeh/bokeh_plot_posterior_combinedims.py b/examples/bokeh/bokeh_plot_posterior_combinedims.py
new file mode 100644
index 0000000..1ca0e49
--- /dev/null
+++ b/examples/bokeh/bokeh_plot_posterior_combinedims.py
@@ -0,0 +1,18 @@
+"""
+Posterior Plot (reducing school dimension)
+==========================================
+
+_thumb: .5, .8
+"""
+import arviz as az
+
+data = az.load_arviz_data("centered_eight")
+
+coords = {"school": ["Choate", "Mt. Hermon", "Deerfield"]}
+ax = az.plot_posterior(
+    data,
+    var_names=["mu", "theta"],
+    combine_dims={"school"},
+    coords=coords,
+    backend="bokeh",
+)
diff --git a/examples/matplotlib/mpl_plot_posterior_combinedims.py b/examples/matplotlib/mpl_plot_posterior_combinedims.py
new file mode 100644
index 0000000..3c34638
--- /dev/null
+++ b/examples/matplotlib/mpl_plot_posterior_combinedims.py
@@ -0,0 +1,17 @@
+"""
+Posterior Plot (reducing school dimension)
+==========================================
+
+_thumb: .5, .8
+"""
+import matplotlib.pyplot as plt
+import arviz as az
+
+az.style.use("arviz-darkgrid")
+
+data = az.load_arviz_data("centered_eight")
+
+coords = {"school": ["Choate", "Mt. Hermon", "Deerfield"]}
+az.plot_posterior(data, var_names=["mu", "theta"], combine_dims={"school"}, coords=coords)
+
+plt.show()


Test patch:
diff --git a/arviz/tests/base_tests/test_plots_bokeh.py b/arviz/tests/base_tests/test_plots_bokeh.py
index c1fb847..329e898 100644
--- a/arviz/tests/base_tests/test_plots_bokeh.py
+++ b/arviz/tests/base_tests/test_plots_bokeh.py
@@ -408,7 +408,7 @@ def test_plot_energy_bad(models):
     [
         {},
         {"var_names": ["theta"], "relative": True, "color": "r"},
-        {"coords": {"theta_dim_0": slice(4)}, "n_points": 10},
+        {"coords": {"school": slice(4)}, "n_points": 10},
         {"min_ess": 600, "hline_kwargs": {"color": "r"}},
     ],
 )
@@ -710,7 +710,7 @@ def test_plot_loo_pit_label(models, args):
         {"var_names": ["theta"], "color": "r"},
         {"rug": True, "rug_kwargs": {"color": "r"}},
         {"errorbar": True, "rug": True, "rug_kind": "max_depth"},
-        {"errorbar": True, "coords": {"theta_dim_0": slice(4)}, "n_points": 10},
+        {"errorbar": True, "coords": {"school": slice(4)}, "n_points": 10},
         {"extra_methods": True, "rug": True},
         {"extra_methods": True, "extra_kwargs": {"ls": ":"}, "text_kwargs": {"x": 0, "ha": "left"}},
     ],
@@ -748,7 +748,7 @@ def test_plot_mcse_no_divergences(models):
 @pytest.mark.parametrize(
     "kwargs",
     [
-        {"var_names": "theta", "divergences": True, "coords": {"theta_dim_0": [0, 1]}},
+        {"var_names": "theta", "divergences": True, "coords": {"school": [0, 1]}},
         {"divergences": True, "var_names": ["theta", "mu"]},
         {"kind": "kde", "var_names": ["theta"]},
         {"kind": "hexbin", "var_names": ["theta"]},
@@ -756,7 +756,7 @@ def test_plot_mcse_no_divergences(models):
         {
             "kind": "hexbin",
             "var_names": ["theta"],
-            "coords": {"theta_dim_0": [0, 1]},
+            "coords": {"school": [0, 1]},
             "textsize": 20,
         },
         {
@@ -1022,7 +1022,7 @@ def test_plot_posterior_skipna():
     [
         {},
         {"var_names": "mu"},
-        {"var_names": ("mu", "tau"), "coords": {"theta_dim_0": [0, 1]}},
+        {"var_names": ("mu", "tau"), "coords": {"school": [0, 1]}},
         {"var_names": "mu", "ref_line": True},
         {
             "var_names": "mu",
diff --git a/arviz/tests/base_tests/test_plots_matplotlib.py b/arviz/tests/base_tests/test_plots_matplotlib.py
index b596667..1fa7af4 100644
--- a/arviz/tests/base_tests/test_plots_matplotlib.py
+++ b/arviz/tests/base_tests/test_plots_matplotlib.py
@@ -85,6 +85,16 @@ def discrete_model():
     return {"x": np.random.randint(10, size=100), "y": np.random.randint(10, size=100)}
 
 
+@pytest.fixture(scope="module")
+def discrete_multidim_model():
+    """Simple fixture for random discrete model"""
+    idata = from_dict(
+        {"x": np.random.randint(10, size=(2, 50, 3)), "y": np.random.randint(10, size=(2, 50))},
+        dims={"x": ["school"]},
+    )
+    return idata
+
+
 @pytest.fixture(scope="module")
 def continuous_model():
     """Simple fixture for random continuous model"""
@@ -167,6 +177,11 @@ def test_plot_density_bad_kwargs(models):
         plot_density(obj, filter_vars="bad_value")
 
 
+def test_plot_density_discrete_combinedims(discrete_model):
+    axes = plot_density(discrete_model, combine_dims={"school"}, shade=0.9)
+    assert axes.size == 2
+
+
 @pytest.mark.parametrize(
     "kwargs",
     [
@@ -270,7 +285,7 @@ def test_plot_trace_futurewarning(models, prop):
     [
         ({}, 1),
         ({"var_names": "mu", "transform": lambda x: x + 1}, 1),
-        ({"var_names": "mu", "rope": (-1, 1)}, 1),
+        ({"var_names": "mu", "rope": (-1, 1), "combine_dims": {"school"}}, 1),
         ({"r_hat": True, "quartiles": False}, 2),
         ({"var_names": ["mu"], "colors": "C0", "ess": True, "combined": True}, 2),
         (
@@ -509,7 +524,7 @@ def test_plot_kde_inference_data(models):
         {
             "var_names": "theta",
             "divergences": True,
-            "coords": {"theta_dim_0": [0, 1]},
+            "coords": {"school": [0, 1]},
             "scatter_kwargs": {"marker": "x"},
             "divergences_kwargs": {"marker": "*", "c": "C0"},
         },
@@ -525,7 +540,7 @@ def test_plot_kde_inference_data(models):
         {
             "kind": "hexbin",
             "var_names": ["theta"],
-            "coords": {"theta_dim_0": [0, 1]},
+            "coords": {"school": [0, 1]},
             "colorbar": True,
             "hexbin_kwargs": {"cmap": "viridis"},
             "textsize": 20,
@@ -582,6 +597,17 @@ def test_plot_pair_overlaid(models, kwargs):
     assert ax.shape
 
 
+@pytest.mark.parametrize("marginals", [True, False])
+def test_plot_pair_combinedims(models, marginals):
+    ax = plot_pair(
+        models.model_1, var_names=["eta", "theta"], combine_dims={"school"}, marginals=marginals
+    )
+    if marginals:
+        assert ax.shape == (2, 2)
+    else:
+        assert not isinstance(ax, np.ndarray)
+
+
 @pytest.mark.parametrize("marginals", [True, False])
 @pytest.mark.parametrize("max_subplots", [True, False])
 def test_plot_pair_shapes(marginals, max_subplots):
@@ -873,6 +899,30 @@ def test_plot_violin_discrete(discrete_model):
     assert axes.shape
 
 
+@pytest.mark.parametrize("var_names", (None, "mu", ["mu", "tau"]))
+def test_plot_violin_combinedims(models, var_names):
+    axes = plot_violin(models.model_1, var_names=var_names, combine_dims={"school"})
+    assert axes.shape
+
+
+def test_plot_violin_ax_combinedims(models):
+    _, ax = plt.subplots(1)
+    axes = plot_violin(models.model_1, var_names="mu", combine_dims={"school"}, ax=ax)
+    assert axes.shape
+
+
+def test_plot_violin_layout_combinedims(models):
+    axes = plot_violin(
+        models.model_1, var_names=["mu", "tau"], combine_dims={"school"}, sharey=False
+    )
+    assert axes.shape
+
+
+def test_plot_violin_discrete_combinedims(discrete_model):
+    axes = plot_violin(discrete_model, combine_dims={"school"})
+    assert axes.shape
+
+
 def test_plot_autocorr_short_chain():
     """Check that logic for small chain defaulting doesn't cause exception"""
     chain = np.arange(10)
@@ -908,7 +958,7 @@ def test_plot_autocorr_var_names(models, var_names):
     [
         {},
         {"var_names": "mu"},
-        {"var_names": ("mu", "tau"), "coords": {"theta_dim_0": [0, 1]}},
+        {"var_names": ("mu", "tau"), "coords": {"school": [0, 1]}},
         {"var_names": "mu", "ref_line": True},
         {
             "var_names": "mu",
@@ -997,6 +1047,41 @@ def test_plot_posterior_skipna():
         plot_posterior({"a": sample}, skipna=False)
 
 
+@pytest.mark.parametrize("kwargs", [{"var_names": ["mu", "theta"]}])
+def test_plot_posterior_combinedims(models, kwargs):
+    axes = plot_posterior(models.model_1, combine_dims={"school"}, **kwargs)
+    if isinstance(kwargs.get("var_names"), str):
+        assert not isinstance(axes, np.ndarray)
+    else:
+        assert axes.shape
+
+
+@pytest.mark.parametrize("kwargs", [{}, {"point_estimate": "mode"}, {"bins": None, "kind": "hist"}])
+def test_plot_posterior_discrete_combinedims(discrete_multidim_model, kwargs):
+    axes = plot_posterior(discrete_multidim_model, combine_dims={"school"}, **kwargs)
+    assert axes.size == 2
+
+
+@pytest.mark.parametrize("point_estimate", ("mode", "mean", "median"))
+def test_plot_posterior_point_estimates_combinedims(models, point_estimate):
+    axes = plot_posterior(
+        models.model_1,
+        var_names=("mu", "tau"),
+        combine_dims={"school"},
+        point_estimate=point_estimate,
+    )
+    assert axes.size == 2
+
+
+def test_plot_posterior_skipna_combinedims():
+    idata = load_arviz_data("centered_eight")
+    idata.posterior["theta"].loc[dict(school="Deerfield")] = np.nan
+    with pytest.raises(ValueError):
+        plot_posterior(idata, var_names="theta", combine_dims={"school"}, skipna=False)
+    ax = plot_posterior(idata, var_names="theta", combine_dims={"school"}, skipna=True)
+    assert not isinstance(ax, np.ndarray)
+
+
 @pytest.mark.parametrize(
     "kwargs", [{"insample_dev": False}, {"plot_standard_error": False}, {"plot_ic_diff": False}]
 )
@@ -1271,7 +1356,7 @@ def test_plot_khat_bad_input(models):
     [
         {},
         {"var_names": ["theta"], "relative": True, "color": "r"},
-        {"coords": {"theta_dim_0": slice(4)}, "n_points": 10},
+        {"coords": {"school": slice(4)}, "n_points": 10},
         {"min_ess": 600, "hline_kwargs": {"color": "r"}},
     ],
 )
@@ -1369,7 +1454,7 @@ def test_plot_loo_pit_incompatible_args(models):
         {"var_names": ["theta"], "color": "r"},
         {"rug": True, "rug_kwargs": {"color": "r"}},
         {"errorbar": True, "rug": True, "rug_kind": "max_depth"},
-        {"errorbar": True, "coords": {"theta_dim_0": slice(4)}, "n_points": 10},
+        {"errorbar": True, "coords": {"school": slice(4)}, "n_points": 10},
         {"extra_methods": True, "rug": True},
         {"extra_methods": True, "extra_kwargs": {"ls": ":"}, "text_kwargs": {"x": 0, "ha": "left"}},
     ],
@@ -1408,7 +1493,7 @@ def test_plot_mcse_no_divergences(models):
     [
         {},
         {"var_names": ["theta"]},
-        {"var_names": ["theta"], "coords": {"theta_dim_0": [0, 1]}},
+        {"var_names": ["theta"], "coords": {"school": [0, 1]}},
         {"var_names": ["eta"], "posterior_kwargs": {"rug": True, "rug_kwargs": {"color": "r"}}},
         {"var_names": ["mu"], "prior_kwargs": {"fill_kwargs": {"alpha": 0.5}}},
         {
@@ -1438,6 +1523,27 @@ def test_plot_dist_comparison_different_vars():
     assert np.all(ax)
 
 
+def test_plot_dist_comparison_combinedims(models):
+    idata = models.model_1
+    ax = plot_dist_comparison(idata, combine_dims={"school"})
+    assert np.all(ax)
+
+
+def test_plot_dist_comparison_different_vars_combinedims():
+    data = from_dict(
+        posterior={
+            "x": np.random.randn(4, 100, 30),
+        },
+        prior={"x_hat": np.random.randn(4, 100, 30)},
+        dims={"x": ["3rd_dim"], "x_hat": ["3rd_dim"]},
+    )
+    with pytest.raises(KeyError):
+        plot_dist_comparison(data, var_names="x", combine_dims={"3rd_dim"})
+    ax = plot_dist_comparison(data, var_names=[["x_hat"], ["x"]], combine_dims={"3rd_dim"})
+    assert np.all(ax)
+    assert ax.size == 3
+
+
 @pytest.mark.parametrize(
     "kwargs",
     [
diff --git a/arviz/tests/helpers.py b/arviz/tests/helpers.py
index d26c813..72b76f4 100644
--- a/arviz/tests/helpers.py
+++ b/arviz/tests/helpers.py
@@ -85,7 +85,12 @@ def create_model(seed=10):
         prior_predictive=prior_predictive,
         sample_stats_prior=sample_stats_prior,
         observed_data={"y": data["y"]},
-        dims={"y": ["obs_dim"], "log_likelihood": ["obs_dim"]},
+        dims={
+            "y": ["obs_dim"],
+            "log_likelihood": ["obs_dim"],
+            "theta": ["school"],
+            "eta": ["school"],
+        },
         coords={"obs_dim": range(data["J"])},
     )
     return model