plot_rank change axes to ax, fix plot_trace ax

CHANGELOG.md

@@ -10,6 +10,7 @@
 * Revamped the `hpd` function to make it work with mutidimensional arrays, InferenceData and xarray objects (#1117)
 * Skip test for optional/extra dependencies when not installed (#1113)
 * Add option to display rank plots instead of trace (#1134)
+
 ### Maintenance and fixes
 * Fixed behaviour of `credible_interval=None` in `plot_posterior` (#1115)
 * Fixed hist kind of `plot_dist` with multidimensional input (#1115)
@@ -18,6 +19,8 @@
 * Update Docker building steps (#1127)
 * Updated benchmarks and moved to asv_benchmarks/benchmarks (#1142)
 * Moved `_fast_kde`, `_fast_kde_2d`, `get_bins` and `_sturges_formula` to `numeric_utils` and `get_coords` to `utils` (#1142)
+* Rank plot: rename `axes` argument to `ax` (#1144)
+
 ### Deprecation
 
 ### Documentation

arviz/plots/backends/bokeh/traceplot.py

@@ -86,7 +86,7 @@ def plot_trace(
                 _axes = [bkp.figure(**backend_kwargs), bkp.figure(**backend_kwargs)]
             axes.append(_axes)
 
-    axes = np.array(axes)
+    axes = np.atleast_2d(axes)
 
     cds_data = {}
     cds_var_groups = {}
@@ -333,10 +333,10 @@ def _plot_chains_bokeh(
 
     if kind == "rank_bars":
         value = np.array([item.data[y_name] for item in data.values()])
-        plot_rank(value, kind="bars", axes=ax_trace, backend="bokeh", show=False, **rank_kwargs)
+        plot_rank(value, kind="bars", ax=ax_trace, backend="bokeh", show=False, **rank_kwargs)
     elif kind == "rank_vlines":
         value = np.array([item.data[y_name] for item in data.values()])
-        plot_rank(value, kind="vlines", axes=ax_trace, backend="bokeh", show=False, **rank_kwargs)
+        plot_rank(value, kind="vlines", ax=ax_trace, backend="bokeh", show=False, **rank_kwargs)
 
     if combined:
         rug_kwargs["cds"] = data

arviz/plots/backends/matplotlib/traceplot.py

@@ -137,6 +137,7 @@ def plot_trace(
     if axes is None:
         _, axes = plt.subplots(len(plotters), 2, squeeze=False, figsize=figsize, **backend_kwargs)
 
+    axes = np.atleast_2d(axes)
     # Check the input for lines
     if lines is not None:
         all_var_names = set(plotter[0] for plotter in plotters)
@@ -362,9 +363,9 @@ def _plot_chains_mpl(
             plot_kwargs.pop(chain_prop[0])
 
     if kind == "rank_bars":
-        plot_rank(data=value, kind="bars", axes=axes[idx, 1], **rank_kwargs)
+        plot_rank(data=value, kind="bars", ax=axes[idx, 1], **rank_kwargs)
     elif kind == "rank_vlines":
-        plot_rank(data=value, kind="vlines", axes=axes[idx, 1], **rank_kwargs)
+        plot_rank(data=value, kind="vlines", ax=axes[idx, 1], **rank_kwargs)
 
     if combined:
         plot_kwargs[chain_prop[0]] = chain_prop[1][-1]

arviz/plots/rankplot.py

@@ -26,7 +26,7 @@ def plot_rank(
     ref_line=True,
     labels=True,
     figsize=None,
-    axes=None,
+    ax=None,
     backend=None,
     backend_kwargs=None,
     show=None,
@@ -136,7 +136,7 @@ def plot_rank(
         bins = _sturges_formula(posterior_data, mult=2)
 
     rows, cols = default_grid(length_plotters)
-    if axes is None:
+    if ax is None:
         figsize, ax_labelsize, titlesize, _, _, _ = _scale_fig_size(
             figsize, None, rows=rows, cols=cols
         )
@@ -155,7 +155,7 @@ def plot_rank(
         colors = [colors] * chains
 
     rankplot_kwargs = dict(
-        axes=axes,
+        axes=ax,
         length_plotters=length_plotters,
         rows=rows,
         cols=cols,