We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
plot_trace assumes specific dimensions ordering for divergences
**Describe the bug**
If an `InferenceData` contains a `diverging` variable in `sample_stats` that uses the dimension ordering `(draw, chain)` instead of `(chain, draw)`, then `plot_trace` fails.

**To Reproduce**
```python
>>> import arviz as az
>>> import numpy as np
>>> idata = az.from_dict(
...     dict(x=np.random.normal(size=(4, 100, 2))),
...     sample_stats=dict(diverging=np.random.randint(0, 1, size=(4, 100), dtype=bool))
... )
>>> az.plot_trace(idata) # fine
array([[<AxesSubplot:title={'center':'x'}>,
        <AxesSubplot:title={'center':'x'}>]], dtype=object)
>>> idata.posterior['x'] = idata.posterior.x.transpose()
>>> idata.sample_stats['diverging'] = idata.sample_stats.diverging.transpose()
>>> az.plot_trace(idata) # error
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/sethaxen/projects/arviz/arviz/plots/traceplot.py", line 263, in plot_trace
    axes = plot(**trace_plot_args)
  File "/home/sethaxen/projects/arviz/arviz/plots/backends/matplotlib/traceplot.py", line 347, in plot_trace
    div_draws = data.draw.values[chain_divs]
IndexError: boolean index did not match indexed array along dimension 0; dimension is 100 but corresponding boolean dimension is 4
>>> idata = az.InferenceData(posterior=idata.posterior)
>>> az.plot_trace(idata) # fine without diverging
array([[<AxesSubplot:title={'center':'x'}>,
        <AxesSubplot:title={'center':'x'}>]], dtype=object)
```

**Expected behavior**
I expected all calls to `plot_trace` to work above.

**Additional context**
These checks were performed on main.


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `arviz/plots/traceplot.py`
Content:
```
1 """Plot kde or histograms and values from MCMC samples."""
2 import warnings
3 from typing import Any, Callable, List, Mapping, Optional, Tuple, Union, Sequence
4 
5 from ..data import CoordSpec, InferenceData, convert_to_dataset
6 from ..labels import BaseLabeller
7 from ..rcparams import rcParams
8 from ..sel_utils import xarray_var_iter
9 from ..utils import _var_names, get_coords
10 from .plot_utils import KwargSpec, get_plotting_function
11 
12 
13 def plot_trace(
14     data: InferenceData,
15     var_names: Optional[Sequence[str]] = None,
16     filter_vars: Optional[str] = None,
17     transform: Optional[Callable] = None,
18     coords: Optional[CoordSpec] = None,
19     divergences: Optional[str] = "auto",
20     kind: Optional[str] = "trace",
21     figsize: Optional[Tuple[float, float]] = None,
22     rug: bool = False,
23     lines: Optional[List[Tuple[str, CoordSpec, Any]]] = None,
24     circ_var_names: Optional[List[str]] = None,
25     circ_var_units: str = "radians",
26     compact: bool = True,
27     compact_prop: Optional[Union[str, Mapping[str, Any]]] = None,
28     combined: bool = False,
29     chain_prop: Optional[Union[str, Mapping[str, Any]]] = None,
30     legend: bool = False,
31     plot_kwargs: Optional[KwargSpec] = None,
32     fill_kwargs: Optional[KwargSpec] = None,
33     rug_kwargs: Optional[KwargSpec] = None,
34     hist_kwargs: Optional[KwargSpec] = None,
35     trace_kwargs: Optional[KwargSpec] = None,
36     rank_kwargs: Optional[KwargSpec] = None,
37     labeller=None,
38     axes=None,
39     backend: Optional[str] = None,
40     backend_config: Optional[KwargSpec] = None,
41     backend_kwargs: Optional[KwargSpec] = None,
42     show: Optional[bool] = None,
43 ):
44     """Plot distribution (histogram or kernel density estimates) and sampled values or rank plot.
45 
46     If `divergences` data is available in `sample_stats`, will plot the location of divergences as
47     dashed vertical lines.
48 
49     Parameters
50     ----------
51     data: obj
52         Any object that can be converted to an :class:`arviz.InferenceData` object
53         Refer to documentation of :func:`arviz.convert_to_dataset` for details
54     var_names: str or list of str, optional
55         One or more variables to be plotted. Prefix the variables by ``~`` when you want
56         to exclude them from the plot.
57     filter_vars: {None, "like", "regex"}, optional, default=None
58         If `None` (default), interpret var_names as the real variables names. If "like",
59         interpret var_names as substrings of the real variables names. If "regex",
60         interpret var_names as regular expressions on the real variables names. A la
61         ``pandas.filter``.
62     coords: dict of {str: slice or array_like}, optional
63         Coordinates of var_names to be plotted. Passed to :meth:`xarray.Dataset.sel`
64     divergences: {"bottom", "top", None}, optional
65         Plot location of divergences on the traceplots.
66     kind: {"trace", "rank_bars", "rank_vlines"}, optional
67         Choose between plotting sampled values per iteration and rank plots.
68     transform: callable, optional
69         Function to transform data (defaults to None i.e.the identity function)
70     figsize: tuple of (float, float), optional
71         If None, size is (12, variables * 2)
72     rug: bool, optional
73         If True adds a rugplot of samples. Defaults to False. Ignored for 2D KDE.
74         Only affects continuous variables.
75     lines: list of tuple of (str, dict, array_like), optional
76         List of (var_name, {'coord': selection}, [line, positions]) to be overplotted as
77         vertical lines on the density and horizontal lines on the trace.
78     circ_var_names : str or list of str, optional
79         List of circular variables to account for when plotting KDE.
80     circ_var_units : str
81         Whether the variables in ``circ_var_names`` are in "degrees" or "radians".
82     compact: bool, optional
83         Plot multidimensional variables in a single plot.
84     compact_prop: str or dict {str: array_like}, optional
85         Tuple containing the property name and the property values to distinguish different
86         dimensions with compact=True
87     combined: bool, optional
88         Flag for combining multiple chains into a single line. If False (default), chains will be
89         plotted separately.
90     chain_prop: str or dict {str: array_like}, optional
91         Tuple containing the property name and the property values to distinguish different chains
92     legend: bool, optional
93         Add a legend to the figure with the chain color code.
94     plot_kwargs, fill_kwargs, rug_kwargs, hist_kwargs: dict, optional
95         Extra keyword arguments passed to :func:`arviz.plot_dist`. Only affects continuous
96         variables.
97     trace_kwargs: dict, optional
98         Extra keyword arguments passed to :meth:`matplotlib.axes.Axes.plot`
99     labeller : labeller instance, optional
100         Class providing the method ``make_label_vert`` to generate the labels in the plot titles.
101         Read the :ref:`label_guide` for more details and usage examples.
102     rank_kwargs : dict, optional
103         Extra keyword arguments passed to :func:`arviz.plot_rank`
104     axes: axes, optional
105         Matplotlib axes or bokeh figures.
106     backend: {"matplotlib", "bokeh"}, optional
107         Select plotting backend.
108     backend_config: dict, optional
109         Currently specifies the bounds to use for bokeh axes. Defaults to value set in rcParams.
110     backend_kwargs: dict, optional
111         These are kwargs specific to the backend being used, passed to
112         :func:`matplotlib.pyplot.subplots` or
113         :func:`bokeh.plotting.figure`.
114     show: bool, optional
115         Call backend show function.
116 
117     Returns
118     -------
119     axes: matplotlib axes or bokeh figures
120 
121     See Also
122     --------
123     plot_rank : Plot rank order statistics of chains.
124 
125     Examples
126     --------
127     Plot a subset variables and select them with partial naming
128 
129     .. plot::
130         :context: close-figs
131 
132         >>> import arviz as az
133         >>> data = az.load_arviz_data('non_centered_eight')
134         >>> coords = {'school': ['Choate', 'Lawrenceville']}
135         >>> az.plot_trace(data, var_names=('theta'), filter_vars="like", coords=coords)
136 
137     Show all dimensions of multidimensional variables in the same plot
138 
139     .. plot::
140         :context: close-figs
141 
142         >>> az.plot_trace(data, compact=True)
143 
144     Display a rank plot instead of trace
145 
146     .. plot::
147         :context: close-figs
148 
149         >>> az.plot_trace(data, var_names=["mu", "tau"], kind="rank_bars")
150 
151     Combine all chains into one distribution and select variables with regular expressions
152 
153     .. plot::
154         :context: close-figs
155 
156         >>> az.plot_trace(
157         >>>     data, var_names=('^theta'), filter_vars="regex", coords=coords, combined=True
158         >>> )
159 
160 
161     Plot reference lines against distribution and trace
162 
163     .. plot::
164         :context: close-figs
165 
166         >>> lines = (('theta_t',{'school': "Choate"}, [-1]),)
167         >>> az.plot_trace(data, var_names=('theta_t', 'theta'), coords=coords, lines=lines)
168 
169     """
170     if kind not in {"trace", "rank_vlines", "rank_bars"}:
171         raise ValueError("The value of kind must be either trace, rank_vlines or rank_bars.")
172 
173     if divergences == "auto":
174         divergences = "top" if rug else "bottom"
175     if divergences:
176         try:
177             divergence_data = convert_to_dataset(data, group="sample_stats").diverging
178         except (ValueError, AttributeError):  # No sample_stats, or no `.diverging`
179             divergences = None
180 
181     if coords is None:
182         coords = {}
183 
184     if labeller is None:
185         labeller = BaseLabeller()
186 
187     if divergences:
188         divergence_data = get_coords(
189             divergence_data, {k: v for k, v in coords.items() if k in ("chain", "draw")}
190         )
191     else:
192         divergence_data = False
193 
194     coords_data = get_coords(convert_to_dataset(data, group="posterior"), coords)
195 
196     if transform is not None:
197         coords_data = transform(coords_data)
198 
199     var_names = _var_names(var_names, coords_data, filter_vars)
200 
201     skip_dims = set(coords_data.dims) - {"chain", "draw"} if compact else set()
202 
203     plotters = list(
204         xarray_var_iter(
205             coords_data,
206             var_names=var_names,
207             combined=True,
208             skip_dims=skip_dims,
209             dim_order=["chain", "draw"],
210         )
211     )
212     max_plots = rcParams["plot.max_subplots"]
213     max_plots = len(plotters) if max_plots is None else max(max_plots // 2, 1)
214     if len(plotters) > max_plots:
215         warnings.warn(
216             "rcParams['plot.max_subplots'] ({max_plots}) is smaller than the number "
217             "of variables to plot ({len_plotters}), generating only {max_plots} "
218             "plots".format(max_plots=max_plots, len_plotters=len(plotters)),
219             UserWarning,
220         )
221         plotters = plotters[:max_plots]
222 
223     # TODO: Check if this can be further simplified
224     trace_plot_args = dict(
225         # User Kwargs
226         data=coords_data,
227         var_names=var_names,
228         # coords = coords,
229         divergences=divergences,
230         kind=kind,
231         figsize=figsize,
232         rug=rug,
233         lines=lines,
234         circ_var_names=circ_var_names,
235         circ_var_units=circ_var_units,
236         plot_kwargs=plot_kwargs,
237         fill_kwargs=fill_kwargs,
238         rug_kwargs=rug_kwargs,
239         hist_kwargs=hist_kwargs,
240         trace_kwargs=trace_kwargs,
241         rank_kwargs=rank_kwargs,
242         compact=compact,
243         compact_prop=compact_prop,
244         combined=combined,
245         chain_prop=chain_prop,
246         legend=legend,
247         labeller=labeller,
248         # Generated kwargs
249         divergence_data=divergence_data,
250         # skip_dims=skip_dims,
251         plotters=plotters,
252         axes=axes,
253         backend_config=backend_config,
254         backend_kwargs=backend_kwargs,
255         show=show,
256     )
257 
258     if backend is None:
259         backend = rcParams["plot.backend"]
260     backend = backend.lower()
261 
262     plot = get_plotting_function("plot_trace", "traceplot", backend)
263     axes = plot(**trace_plot_args)
264 
265     return axes
266
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/arviz/plots/traceplot.py b/arviz/plots/traceplot.py
--- a/arviz/plots/traceplot.py
+++ b/arviz/plots/traceplot.py
@@ -174,7 +174,9 @@
         divergences = "top" if rug else "bottom"
     if divergences:
         try:
-            divergence_data = convert_to_dataset(data, group="sample_stats").diverging
+            divergence_data = convert_to_dataset(data, group="sample_stats").diverging.transpose(
+                "chain", "draw"
+            )
         except (ValueError, AttributeError):  # No sample_stats, or no `.diverging`
             divergences = None

{"golden_diff": "diff --git a/arviz/plots/traceplot.py b/arviz/plots/traceplot.py\n--- a/arviz/plots/traceplot.py\n+++ b/arviz/plots/traceplot.py\n@@ -174,7 +174,9 @@\n         divergences = \"top\" if rug else \"bottom\"\n     if divergences:\n         try:\n-            divergence_data = convert_to_dataset(data, group=\"sample_stats\").diverging\n+            divergence_data = convert_to_dataset(data, group=\"sample_stats\").diverging.transpose(\n+                \"chain\", \"draw\"\n+            )\n         except (ValueError, AttributeError):  # No sample_stats, or no `.diverging`\n             divergences = None\n", "issue": "plot_trace assumes specific dimensions ordering for divergences\n**Describe the bug**\r\nIf an `InferenceData` contains a `diverging` variable in `sample_stats` that uses the dimension ordering `(draw, chain)` instead of `(chain, draw)`, then `plot_trace` fails.\r\n\r\n**To Reproduce**\r\n```python\r\n>>> import arviz as az\r\n>>> import numpy as np\r\n>>> idata = az.from_dict(\r\n...     dict(x=np.random.normal(size=(4, 100, 2))),\r\n...     sample_stats=dict(diverging=np.random.randint(0, 1, size=(4, 100), dtype=bool))\r\n... )\r\n>>> az.plot_trace(idata) # fine\r\narray([[<AxesSubplot:title={'center':'x'}>,\r\n        <AxesSubplot:title={'center':'x'}>]], dtype=object)\r\n>>> idata.posterior['x'] = idata.posterior.x.transpose()\r\n>>> idata.sample_stats['diverging'] = idata.sample_stats.diverging.transpose()\r\n>>> az.plot_trace(idata) # error\r\nTraceback (most recent call last):\r\n  File \"<stdin>\", line 1, in <module>\r\n  File \"/home/sethaxen/projects/arviz/arviz/plots/traceplot.py\", line 263, in plot_trace\r\n    axes = plot(**trace_plot_args)\r\n  File \"/home/sethaxen/projects/arviz/arviz/plots/backends/matplotlib/traceplot.py\", line 347, in plot_trace\r\n    div_draws = data.draw.values[chain_divs]\r\nIndexError: boolean index did not match indexed array along dimension 0; dimension is 100 but corresponding boolean dimension is 4\r\n>>> idata = az.InferenceData(posterior=idata.posterior)\r\n>>> az.plot_trace(idata) # fine without diverging\r\narray([[<AxesSubplot:title={'center':'x'}>,\r\n        <AxesSubplot:title={'center':'x'}>]], dtype=object)\r\n```\r\n\r\n**Expected behavior**\r\nI expected all calls to `plot_trace` to work above.\r\n\r\n**Additional context**\r\nThese checks were performed on main.\r\n\n", "before_files": [{"content": "\"\"\"Plot kde or histograms and values from MCMC samples.\"\"\"\nimport warnings\nfrom typing import Any, Callable, List, Mapping, Optional, Tuple, Union, Sequence\n\nfrom ..data import CoordSpec, InferenceData, convert_to_dataset\nfrom ..labels import BaseLabeller\nfrom ..rcparams import rcParams\nfrom ..sel_utils import xarray_var_iter\nfrom ..utils import _var_names, get_coords\nfrom .plot_utils import KwargSpec, get_plotting_function\n\n\ndef plot_trace(\n    data: InferenceData,\n    var_names: Optional[Sequence[str]] = None,\n    filter_vars: Optional[str] = None,\n    transform: Optional[Callable] = None,\n    coords: Optional[CoordSpec] = None,\n    divergences: Optional[str] = \"auto\",\n    kind: Optional[str] = \"trace\",\n    figsize: Optional[Tuple[float, float]] = None,\n    rug: bool = False,\n    lines: Optional[List[Tuple[str, CoordSpec, Any]]] = None,\n    circ_var_names: Optional[List[str]] = None,\n    circ_var_units: str = \"radians\",\n    compact: bool = True,\n    compact_prop: Optional[Union[str, Mapping[str, Any]]] = None,\n    combined: bool = False,\n    chain_prop: Optional[Union[str, Mapping[str, Any]]] = None,\n    legend: bool = False,\n    plot_kwargs: Optional[KwargSpec] = None,\n    fill_kwargs: Optional[KwargSpec] = None,\n    rug_kwargs: Optional[KwargSpec] = None,\n    hist_kwargs: Optional[KwargSpec] = None,\n    trace_kwargs: Optional[KwargSpec] = None,\n    rank_kwargs: Optional[KwargSpec] = None,\n    labeller=None,\n    axes=None,\n    backend: Optional[str] = None,\n    backend_config: Optional[KwargSpec] = None,\n    backend_kwargs: Optional[KwargSpec] = None,\n    show: Optional[bool] = None,\n):\n    \"\"\"Plot distribution (histogram or kernel density estimates) and sampled values or rank plot.\n\n    If `divergences` data is available in `sample_stats`, will plot the location of divergences as\n    dashed vertical lines.\n\n    Parameters\n    ----------\n    data: obj\n        Any object that can be converted to an :class:`arviz.InferenceData` object\n        Refer to documentation of :func:`arviz.convert_to_dataset` for details\n    var_names: str or list of str, optional\n        One or more variables to be plotted. Prefix the variables by ``~`` when you want\n        to exclude them from the plot.\n    filter_vars: {None, \"like\", \"regex\"}, optional, default=None\n        If `None` (default), interpret var_names as the real variables names. If \"like\",\n        interpret var_names as substrings of the real variables names. If \"regex\",\n        interpret var_names as regular expressions on the real variables names. A la\n        ``pandas.filter``.\n    coords: dict of {str: slice or array_like}, optional\n        Coordinates of var_names to be plotted. Passed to :meth:`xarray.Dataset.sel`\n    divergences: {\"bottom\", \"top\", None}, optional\n        Plot location of divergences on the traceplots.\n    kind: {\"trace\", \"rank_bars\", \"rank_vlines\"}, optional\n        Choose between plotting sampled values per iteration and rank plots.\n    transform: callable, optional\n        Function to transform data (defaults to None i.e.the identity function)\n    figsize: tuple of (float, float), optional\n        If None, size is (12, variables * 2)\n    rug: bool, optional\n        If True adds a rugplot of samples. Defaults to False. Ignored for 2D KDE.\n        Only affects continuous variables.\n    lines: list of tuple of (str, dict, array_like), optional\n        List of (var_name, {'coord': selection}, [line, positions]) to be overplotted as\n        vertical lines on the density and horizontal lines on the trace.\n    circ_var_names : str or list of str, optional\n        List of circular variables to account for when plotting KDE.\n    circ_var_units : str\n        Whether the variables in ``circ_var_names`` are in \"degrees\" or \"radians\".\n    compact: bool, optional\n        Plot multidimensional variables in a single plot.\n    compact_prop: str or dict {str: array_like}, optional\n        Tuple containing the property name and the property values to distinguish different\n        dimensions with compact=True\n    combined: bool, optional\n        Flag for combining multiple chains into a single line. If False (default), chains will be\n        plotted separately.\n    chain_prop: str or dict {str: array_like}, optional\n        Tuple containing the property name and the property values to distinguish different chains\n    legend: bool, optional\n        Add a legend to the figure with the chain color code.\n    plot_kwargs, fill_kwargs, rug_kwargs, hist_kwargs: dict, optional\n        Extra keyword arguments passed to :func:`arviz.plot_dist`. Only affects continuous\n        variables.\n    trace_kwargs: dict, optional\n        Extra keyword arguments passed to :meth:`matplotlib.axes.Axes.plot`\n    labeller : labeller instance, optional\n        Class providing the method ``make_label_vert`` to generate the labels in the plot titles.\n        Read the :ref:`label_guide` for more details and usage examples.\n    rank_kwargs : dict, optional\n        Extra keyword arguments passed to :func:`arviz.plot_rank`\n    axes: axes, optional\n        Matplotlib axes or bokeh figures.\n    backend: {\"matplotlib\", \"bokeh\"}, optional\n        Select plotting backend.\n    backend_config: dict, optional\n        Currently specifies the bounds to use for bokeh axes. Defaults to value set in rcParams.\n    backend_kwargs: dict, optional\n        These are kwargs specific to the backend being used, passed to\n        :func:`matplotlib.pyplot.subplots` or\n        :func:`bokeh.plotting.figure`.\n    show: bool, optional\n        Call backend show function.\n\n    Returns\n    -------\n    axes: matplotlib axes or bokeh figures\n\n    See Also\n    --------\n    plot_rank : Plot rank order statistics of chains.\n\n    Examples\n    --------\n    Plot a subset variables and select them with partial naming\n\n    .. plot::\n        :context: close-figs\n\n        >>> import arviz as az\n        >>> data = az.load_arviz_data('non_centered_eight')\n        >>> coords = {'school': ['Choate', 'Lawrenceville']}\n        >>> az.plot_trace(data, var_names=('theta'), filter_vars=\"like\", coords=coords)\n\n    Show all dimensions of multidimensional variables in the same plot\n\n    .. plot::\n        :context: close-figs\n\n        >>> az.plot_trace(data, compact=True)\n\n    Display a rank plot instead of trace\n\n    .. plot::\n        :context: close-figs\n\n        >>> az.plot_trace(data, var_names=[\"mu\", \"tau\"], kind=\"rank_bars\")\n\n    Combine all chains into one distribution and select variables with regular expressions\n\n    .. plot::\n        :context: close-figs\n\n        >>> az.plot_trace(\n        >>>     data, var_names=('^theta'), filter_vars=\"regex\", coords=coords, combined=True\n        >>> )\n\n\n    Plot reference lines against distribution and trace\n\n    .. plot::\n        :context: close-figs\n\n        >>> lines = (('theta_t',{'school': \"Choate\"}, [-1]),)\n        >>> az.plot_trace(data, var_names=('theta_t', 'theta'), coords=coords, lines=lines)\n\n    \"\"\"\n    if kind not in {\"trace\", \"rank_vlines\", \"rank_bars\"}:\n        raise ValueError(\"The value of kind must be either trace, rank_vlines or rank_bars.\")\n\n    if divergences == \"auto\":\n        divergences = \"top\" if rug else \"bottom\"\n    if divergences:\n        try:\n            divergence_data = convert_to_dataset(data, group=\"sample_stats\").diverging\n        except (ValueError, AttributeError):  # No sample_stats, or no `.diverging`\n            divergences = None\n\n    if coords is None:\n        coords = {}\n\n    if labeller is None:\n        labeller = BaseLabeller()\n\n    if divergences:\n        divergence_data = get_coords(\n            divergence_data, {k: v for k, v in coords.items() if k in (\"chain\", \"draw\")}\n        )\n    else:\n        divergence_data = False\n\n    coords_data = get_coords(convert_to_dataset(data, group=\"posterior\"), coords)\n\n    if transform is not None:\n        coords_data = transform(coords_data)\n\n    var_names = _var_names(var_names, coords_data, filter_vars)\n\n    skip_dims = set(coords_data.dims) - {\"chain\", \"draw\"} if compact else set()\n\n    plotters = list(\n        xarray_var_iter(\n            coords_data,\n            var_names=var_names,\n            combined=True,\n            skip_dims=skip_dims,\n            dim_order=[\"chain\", \"draw\"],\n        )\n    )\n    max_plots = rcParams[\"plot.max_subplots\"]\n    max_plots = len(plotters) if max_plots is None else max(max_plots // 2, 1)\n    if len(plotters) > max_plots:\n        warnings.warn(\n            \"rcParams['plot.max_subplots'] ({max_plots}) is smaller than the number \"\n            \"of variables to plot ({len_plotters}), generating only {max_plots} \"\n            \"plots\".format(max_plots=max_plots, len_plotters=len(plotters)),\n            UserWarning,\n        )\n        plotters = plotters[:max_plots]\n\n    # TODO: Check if this can be further simplified\n    trace_plot_args = dict(\n        # User Kwargs\n        data=coords_data,\n        var_names=var_names,\n        # coords = coords,\n        divergences=divergences,\n        kind=kind,\n        figsize=figsize,\n        rug=rug,\n        lines=lines,\n        circ_var_names=circ_var_names,\n        circ_var_units=circ_var_units,\n        plot_kwargs=plot_kwargs,\n        fill_kwargs=fill_kwargs,\n        rug_kwargs=rug_kwargs,\n        hist_kwargs=hist_kwargs,\n        trace_kwargs=trace_kwargs,\n        rank_kwargs=rank_kwargs,\n        compact=compact,\n        compact_prop=compact_prop,\n        combined=combined,\n        chain_prop=chain_prop,\n        legend=legend,\n        labeller=labeller,\n        # Generated kwargs\n        divergence_data=divergence_data,\n        # skip_dims=skip_dims,\n        plotters=plotters,\n        axes=axes,\n        backend_config=backend_config,\n        backend_kwargs=backend_kwargs,\n        show=show,\n    )\n\n    if backend is None:\n        backend = rcParams[\"plot.backend\"]\n    backend = backend.lower()\n\n    plot = get_plotting_function(\"plot_trace\", \"traceplot\", backend)\n    axes = plot(**trace_plot_args)\n\n    return axes\n", "path": "arviz/plots/traceplot.py"}], "after_files": [{"content": "\"\"\"Plot kde or histograms and values from MCMC samples.\"\"\"\nimport warnings\nfrom typing import Any, Callable, List, Mapping, Optional, Tuple, Union, Sequence\n\nfrom ..data import CoordSpec, InferenceData, convert_to_dataset\nfrom ..labels import BaseLabeller\nfrom ..rcparams import rcParams\nfrom ..sel_utils import xarray_var_iter\nfrom ..utils import _var_names, get_coords\nfrom .plot_utils import KwargSpec, get_plotting_function\n\n\ndef plot_trace(\n    data: InferenceData,\n    var_names: Optional[Sequence[str]] = None,\n    filter_vars: Optional[str] = None,\n    transform: Optional[Callable] = None,\n    coords: Optional[CoordSpec] = None,\n    divergences: Optional[str] = \"auto\",\n    kind: Optional[str] = \"trace\",\n    figsize: Optional[Tuple[float, float]] = None,\n    rug: bool = False,\n    lines: Optional[List[Tuple[str, CoordSpec, Any]]] = None,\n    circ_var_names: Optional[List[str]] = None,\n    circ_var_units: str = \"radians\",\n    compact: bool = True,\n    compact_prop: Optional[Union[str, Mapping[str, Any]]] = None,\n    combined: bool = False,\n    chain_prop: Optional[Union[str, Mapping[str, Any]]] = None,\n    legend: bool = False,\n    plot_kwargs: Optional[KwargSpec] = None,\n    fill_kwargs: Optional[KwargSpec] = None,\n    rug_kwargs: Optional[KwargSpec] = None,\n    hist_kwargs: Optional[KwargSpec] = None,\n    trace_kwargs: Optional[KwargSpec] = None,\n    rank_kwargs: Optional[KwargSpec] = None,\n    labeller=None,\n    axes=None,\n    backend: Optional[str] = None,\n    backend_config: Optional[KwargSpec] = None,\n    backend_kwargs: Optional[KwargSpec] = None,\n    show: Optional[bool] = None,\n):\n    \"\"\"Plot distribution (histogram or kernel density estimates) and sampled values or rank plot.\n\n    If `divergences` data is available in `sample_stats`, will plot the location of divergences as\n    dashed vertical lines.\n\n    Parameters\n    ----------\n    data: obj\n        Any object that can be converted to an :class:`arviz.InferenceData` object\n        Refer to documentation of :func:`arviz.convert_to_dataset` for details\n    var_names: str or list of str, optional\n        One or more variables to be plotted. Prefix the variables by ``~`` when you want\n        to exclude them from the plot.\n    filter_vars: {None, \"like\", \"regex\"}, optional, default=None\n        If `None` (default), interpret var_names as the real variables names. If \"like\",\n        interpret var_names as substrings of the real variables names. If \"regex\",\n        interpret var_names as regular expressions on the real variables names. A la\n        ``pandas.filter``.\n    coords: dict of {str: slice or array_like}, optional\n        Coordinates of var_names to be plotted. Passed to :meth:`xarray.Dataset.sel`\n    divergences: {\"bottom\", \"top\", None}, optional\n        Plot location of divergences on the traceplots.\n    kind: {\"trace\", \"rank_bars\", \"rank_vlines\"}, optional\n        Choose between plotting sampled values per iteration and rank plots.\n    transform: callable, optional\n        Function to transform data (defaults to None i.e.the identity function)\n    figsize: tuple of (float, float), optional\n        If None, size is (12, variables * 2)\n    rug: bool, optional\n        If True adds a rugplot of samples. Defaults to False. Ignored for 2D KDE.\n        Only affects continuous variables.\n    lines: list of tuple of (str, dict, array_like), optional\n        List of (var_name, {'coord': selection}, [line, positions]) to be overplotted as\n        vertical lines on the density and horizontal lines on the trace.\n    circ_var_names : str or list of str, optional\n        List of circular variables to account for when plotting KDE.\n    circ_var_units : str\n        Whether the variables in ``circ_var_names`` are in \"degrees\" or \"radians\".\n    compact: bool, optional\n        Plot multidimensional variables in a single plot.\n    compact_prop: str or dict {str: array_like}, optional\n        Tuple containing the property name and the property values to distinguish different\n        dimensions with compact=True\n    combined: bool, optional\n        Flag for combining multiple chains into a single line. If False (default), chains will be\n        plotted separately.\n    chain_prop: str or dict {str: array_like}, optional\n        Tuple containing the property name and the property values to distinguish different chains\n    legend: bool, optional\n        Add a legend to the figure with the chain color code.\n    plot_kwargs, fill_kwargs, rug_kwargs, hist_kwargs: dict, optional\n        Extra keyword arguments passed to :func:`arviz.plot_dist`. Only affects continuous\n        variables.\n    trace_kwargs: dict, optional\n        Extra keyword arguments passed to :meth:`matplotlib.axes.Axes.plot`\n    labeller : labeller instance, optional\n        Class providing the method ``make_label_vert`` to generate the labels in the plot titles.\n        Read the :ref:`label_guide` for more details and usage examples.\n    rank_kwargs : dict, optional\n        Extra keyword arguments passed to :func:`arviz.plot_rank`\n    axes: axes, optional\n        Matplotlib axes or bokeh figures.\n    backend: {\"matplotlib\", \"bokeh\"}, optional\n        Select plotting backend.\n    backend_config: dict, optional\n        Currently specifies the bounds to use for bokeh axes. Defaults to value set in rcParams.\n    backend_kwargs: dict, optional\n        These are kwargs specific to the backend being used, passed to\n        :func:`matplotlib.pyplot.subplots` or\n        :func:`bokeh.plotting.figure`.\n    show: bool, optional\n        Call backend show function.\n\n    Returns\n    -------\n    axes: matplotlib axes or bokeh figures\n\n    See Also\n    --------\n    plot_rank : Plot rank order statistics of chains.\n\n    Examples\n    --------\n    Plot a subset variables and select them with partial naming\n\n    .. plot::\n        :context: close-figs\n\n        >>> import arviz as az\n        >>> data = az.load_arviz_data('non_centered_eight')\n        >>> coords = {'school': ['Choate', 'Lawrenceville']}\n        >>> az.plot_trace(data, var_names=('theta'), filter_vars=\"like\", coords=coords)\n\n    Show all dimensions of multidimensional variables in the same plot\n\n    .. plot::\n        :context: close-figs\n\n        >>> az.plot_trace(data, compact=True)\n\n    Display a rank plot instead of trace\n\n    .. plot::\n        :context: close-figs\n\n        >>> az.plot_trace(data, var_names=[\"mu\", \"tau\"], kind=\"rank_bars\")\n\n    Combine all chains into one distribution and select variables with regular expressions\n\n    .. plot::\n        :context: close-figs\n\n        >>> az.plot_trace(\n        >>>     data, var_names=('^theta'), filter_vars=\"regex\", coords=coords, combined=True\n        >>> )\n\n\n    Plot reference lines against distribution and trace\n\n    .. plot::\n        :context: close-figs\n\n        >>> lines = (('theta_t',{'school': \"Choate\"}, [-1]),)\n        >>> az.plot_trace(data, var_names=('theta_t', 'theta'), coords=coords, lines=lines)\n\n    \"\"\"\n    if kind not in {\"trace\", \"rank_vlines\", \"rank_bars\"}:\n        raise ValueError(\"The value of kind must be either trace, rank_vlines or rank_bars.\")\n\n    if divergences == \"auto\":\n        divergences = \"top\" if rug else \"bottom\"\n    if divergences:\n        try:\n            divergence_data = convert_to_dataset(data, group=\"sample_stats\").diverging.transpose(\n                \"chain\", \"draw\"\n            )\n        except (ValueError, AttributeError):  # No sample_stats, or no `.diverging`\n            divergences = None\n\n    if coords is None:\n        coords = {}\n\n    if labeller is None:\n        labeller = BaseLabeller()\n\n    if divergences:\n        divergence_data = get_coords(\n            divergence_data, {k: v for k, v in coords.items() if k in (\"chain\", \"draw\")}\n        )\n    else:\n        divergence_data = False\n\n    coords_data = get_coords(convert_to_dataset(data, group=\"posterior\"), coords)\n\n    if transform is not None:\n        coords_data = transform(coords_data)\n\n    var_names = _var_names(var_names, coords_data, filter_vars)\n\n    skip_dims = set(coords_data.dims) - {\"chain\", \"draw\"} if compact else set()\n\n    plotters = list(\n        xarray_var_iter(\n            coords_data,\n            var_names=var_names,\n            combined=True,\n            skip_dims=skip_dims,\n            dim_order=[\"chain\", \"draw\"],\n        )\n    )\n    max_plots = rcParams[\"plot.max_subplots\"]\n    max_plots = len(plotters) if max_plots is None else max(max_plots // 2, 1)\n    if len(plotters) > max_plots:\n        warnings.warn(\n            \"rcParams['plot.max_subplots'] ({max_plots}) is smaller than the number \"\n            \"of variables to plot ({len_plotters}), generating only {max_plots} \"\n            \"plots\".format(max_plots=max_plots, len_plotters=len(plotters)),\n            UserWarning,\n        )\n        plotters = plotters[:max_plots]\n\n    # TODO: Check if this can be further simplified\n    trace_plot_args = dict(\n        # User Kwargs\n        data=coords_data,\n        var_names=var_names,\n        # coords = coords,\n        divergences=divergences,\n        kind=kind,\n        figsize=figsize,\n        rug=rug,\n        lines=lines,\n        circ_var_names=circ_var_names,\n        circ_var_units=circ_var_units,\n        plot_kwargs=plot_kwargs,\n        fill_kwargs=fill_kwargs,\n        rug_kwargs=rug_kwargs,\n        hist_kwargs=hist_kwargs,\n        trace_kwargs=trace_kwargs,\n        rank_kwargs=rank_kwargs,\n        compact=compact,\n        compact_prop=compact_prop,\n        combined=combined,\n        chain_prop=chain_prop,\n        legend=legend,\n        labeller=labeller,\n        # Generated kwargs\n        divergence_data=divergence_data,\n        # skip_dims=skip_dims,\n        plotters=plotters,\n        axes=axes,\n        backend_config=backend_config,\n        backend_kwargs=backend_kwargs,\n        show=show,\n    )\n\n    if backend is None:\n        backend = rcParams[\"plot.backend\"]\n    backend = backend.lower()\n\n    plot = get_plotting_function(\"plot_trace\", \"traceplot\", backend)\n    axes = plot(**trace_plot_args)\n\n    return axes\n", "path": "arviz/plots/traceplot.py"}]}