Repository: arviz-devs/arviz

Issue:
Geweke diagnostic used inside numba user guide
Numba [user guide](https://arviz-devs.github.io/arviz/user_guide/Numba.html) imports geweke diagnostic even though its removed from arviz (#1545), this results in an `ImportError`.

```python
---------------------------------------------------------------------------
ImportError                               Traceback (most recent call last)
<ipython-input-4-101763bedd0c> in <module>
      4 
      5 from arviz.utils import conditional_jit, Numba
----> 6 from arviz.stats import geweke
      7 from arviz.stats.diagnostics import ks_summary

ImportError: cannot import name 'geweke' from 'arviz.stats' (/oss/libs/arviz/arviz/stats/__init__.py)
```

Discussion:
OriolAbril: an you reexecute the notebook? It looks like it has been a while since it was executed.
MgeeeeK: > an you reexecute the notebook? It looks like it has been a while since it was executed.

Yes, I did run all the cells... results on my machine were not as dramatic, that's why I didn't include those  

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 67d44fd..5c3f30c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -30,6 +30,7 @@
 * Upload updated mypy.ini in ci if mypy copilot fails ([1624](https://github.com/arviz-devs/arviz/pull/1624))
 * Added type checking to raise an error whenever `InferenceData` object is passed using `io_pymc3`'s `trace` argument ([1629](https://github.com/arviz-devs/arviz/pull/1629))
 * Fix `xlabels` in `plot_elpd` ([1601](https://github.com/arviz-devs/arviz/pull/1601))
+* Renamed `sample` dim to `__sample__` when stacking `chain` and `draw` to avoid dimension collision ([1647](https://github.com/arviz-devs/arviz/pull/1647))
 
 ### Deprecation
 * Deprecated `index_origin` and `order` arguments in `az.summary` ([1201](https://github.com/arviz-devs/arviz/pull/1201))
@@ -39,6 +40,7 @@
 * Add "Installation guide" page to the documentation ([1551](https://github.com/arviz-devs/arviz/pull/1551))
 * Improve documentation on experimental `SamplingWrapper` classes ([1582](https://github.com/arviz-devs/arviz/pull/1582))
 * Added example to `plot_hdi` using Inference Data ([1615](https://github.com/arviz-devs/arviz/pull/1615))
+* Removed `geweke` diagnostic from `numba` user guide ([1653](https://github.com/arviz-devs/arviz/pull/1653))
 
 ## v0.11.2 (2021 Feb 21)
 ### New features
diff --git a/arviz/stats/stats.py b/arviz/stats/stats.py
index f3fadb3..5467e34 100644
--- a/arviz/stats/stats.py
+++ b/arviz/stats/stats.py
@@ -648,7 +648,7 @@ def loo(data, pointwise=None, var_name=None, reff=None, scale=None):
     log_likelihood = _get_log_likelihood(inference_data, var_name=var_name)
     pointwise = rcParams["stats.ic_pointwise"] if pointwise is None else pointwise
 
-    log_likelihood = log_likelihood.stack(sample=("chain", "draw"))
+    log_likelihood = log_likelihood.stack(__sample__=("chain", "draw"))
     shape = log_likelihood.shape
     n_samples = shape[-1]
     n_data_points = np.product(shape[:-1])
@@ -692,7 +692,7 @@ def loo(data, pointwise=None, var_name=None, reff=None, scale=None):
         warn_mg = True
 
     ufunc_kwargs = {"n_dims": 1, "ravel": False}
-    kwargs = {"input_core_dims": [["sample"]]}
+    kwargs = {"input_core_dims": [["__sample__"]]}
     loo_lppd_i = scale_value * _wrap_xarray_ufunc(
         _logsumexp, log_weights, ufunc_kwargs=ufunc_kwargs, **kwargs
     )
@@ -752,6 +752,14 @@ def psislw(log_weights, reff=1.0):
     """
     Pareto smoothed importance sampling (PSIS).
 
+    Notes
+    -----
+    If the ``log_weights`` input is an :class:`~xarray.DataArray` with a dimension
+    named ``__sample__`` (recommended) ``psislw`` will interpret this dimension as samples,
+    and all other dimensions as dimensions of the observed data, looping over them to
+    calculate the psislw of each observation. If no ``__sample__`` dimension is present or
+    the input is a numpy array, the last dimension will be interpreted as ``__sample__``.
+
     Parameters
     ----------
     log_weights: array
@@ -778,13 +786,17 @@ def psislw(log_weights, reff=1.0):
 
         In [1]: import arviz as az
            ...: data = az.load_arviz_data("centered_eight")
-           ...: log_likelihood = data.sample_stats.log_likelihood.stack(sample=("chain", "draw"))
+           ...: log_likelihood = data.sample_stats.log_likelihood.stack(
+           ...:     __sample__=("chain", "draw")
+           ...: )
            ...: az.psislw(-log_likelihood, reff=0.8)
 
     """
-    if hasattr(log_weights, "sample"):
-        n_samples = len(log_weights.sample)
-        shape = [size for size, dim in zip(log_weights.shape, log_weights.dims) if dim != "sample"]
+    if hasattr(log_weights, "__sample__"):
+        n_samples = len(log_weights.__sample__)
+        shape = [
+            size for size, dim in zip(log_weights.shape, log_weights.dims) if dim != "__sample__"
+        ]
     else:
         n_samples = log_weights.shape[-1]
         shape = log_weights.shape[:-1]
@@ -799,7 +811,7 @@ def psislw(log_weights, reff=1.0):
     # define kwargs
     func_kwargs = {"cutoff_ind": cutoff_ind, "cutoffmin": cutoffmin, "k_min": k_min, "out": out}
     ufunc_kwargs = {"n_dims": 1, "n_output": 2, "ravel": False, "check_shape": False}
-    kwargs = {"input_core_dims": [["sample"]], "output_core_dims": [["sample"], []]}
+    kwargs = {"input_core_dims": [["__sample__"]], "output_core_dims": [["__sample__"], []]}
     log_weights, pareto_shape = _wrap_xarray_ufunc(
         _psislw,
         log_weights,
@@ -808,7 +820,7 @@ def psislw(log_weights, reff=1.0):
         **kwargs,
     )
     if isinstance(log_weights, xr.DataArray):
-        log_weights = log_weights.rename("log_weights").rename(sample="sample")
+        log_weights = log_weights.rename("log_weights")
     if isinstance(pareto_shape, xr.DataArray):
         pareto_shape = pareto_shape.rename("pareto_shape")
     return log_weights, pareto_shape
@@ -1415,13 +1427,13 @@ def waic(data, pointwise=None, var_name=None, scale=None, dask_kwargs=None):
     else:
         raise TypeError('Valid scale values are "deviance", "log", "negative_log"')
 
-    log_likelihood = log_likelihood.stack(sample=("chain", "draw"))
+    log_likelihood = log_likelihood.stack(__sample__=("chain", "draw"))
     shape = log_likelihood.shape
     n_samples = shape[-1]
     n_data_points = np.product(shape[:-1])
 
     ufunc_kwargs = {"n_dims": 1, "ravel": False}
-    kwargs = {"input_core_dims": [["sample"]]}
+    kwargs = {"input_core_dims": [["__sample__"]]}
     lppd_i = _wrap_xarray_ufunc(
         _logsumexp,
         log_likelihood,
@@ -1431,7 +1443,7 @@ def waic(data, pointwise=None, var_name=None, scale=None, dask_kwargs=None):
         **kwargs,
     )
 
-    vars_lpd = log_likelihood.var(dim="sample")
+    vars_lpd = log_likelihood.var(dim="__sample__")
     warn_mg = False
     if np.any(vars_lpd > 0.4):
         warnings.warn(
@@ -1535,7 +1547,7 @@ def loo_pit(idata=None, *, y=None, y_hat=None, log_weights=None):
 
         In [1]: T = data.observed_data.obs - data.posterior.mu.median(dim=("chain", "draw"))
            ...: T_hat = data.posterior_predictive.obs - data.posterior.mu
-           ...: T_hat = T_hat.stack(sample=("chain", "draw"))
+           ...: T_hat = T_hat.stack(__sample__=("chain", "draw"))
            ...: az.loo_pit(idata=data, y=T**2, y_hat=T_hat**2)
 
     """
@@ -1561,7 +1573,7 @@ def loo_pit(idata=None, *, y=None, y_hat=None, log_weights=None):
         elif not isinstance(y, (np.ndarray, xr.DataArray)):
             raise ValueError(f"y must be of types array, DataArray or str, not {type(y)}")
         if isinstance(y_hat, str):
-            y_hat = idata.posterior_predictive[y_hat].stack(sample=("chain", "draw")).values
+            y_hat = idata.posterior_predictive[y_hat].stack(__sample__=("chain", "draw")).values
         elif not isinstance(y_hat, (np.ndarray, xr.DataArray)):
             raise ValueError(f"y_hat must be of types array, DataArray or str, not {type(y_hat)}")
         if log_weights is None:
@@ -1572,10 +1584,10 @@ def loo_pit(idata=None, *, y=None, y_hat=None, log_weights=None):
                     log_likelihood = _get_log_likelihood(idata)
             else:
                 log_likelihood = _get_log_likelihood(idata)
-            log_likelihood = log_likelihood.stack(sample=("chain", "draw"))
+            log_likelihood = log_likelihood.stack(__sample__=("chain", "draw"))
             posterior = convert_to_dataset(idata, group="posterior")
             n_chains = len(posterior.chain)
-            n_samples = len(log_likelihood.sample)
+            n_samples = len(log_likelihood.__sample__)
             ess_p = ess(posterior, method="mean")
             # this mean is over all data variables
             reff = (
@@ -1608,7 +1620,7 @@ def loo_pit(idata=None, *, y=None, y_hat=None, log_weights=None):
         )
 
     kwargs = {
-        "input_core_dims": [[], ["sample"], ["sample"]],
+        "input_core_dims": [[], ["__sample__"], ["__sample__"]],
         "output_core_dims": [[]],
         "join": "left",
     }
diff --git a/doc/_templates/navbar-version.html b/doc/_templates/navbar-version.html
new file mode 100644
index 0000000..697e752
--- /dev/null
+++ b/doc/_templates/navbar-version.html
@@ -0,0 +1,1 @@
+{{ version }}
diff --git a/doc/source/conf.py b/doc/source/conf.py
index 3c5c3e6..f0ed836 100644
--- a/doc/source/conf.py
+++ b/doc/source/conf.py
@@ -106,6 +106,7 @@ if branch_name == "main":
 else:
     # The short X.Y version.
     version = arviz.__version__
+
 # The full version, including alpha/beta/rc tags.
 release = version
 
@@ -142,8 +143,26 @@ html_theme = "pydata_sphinx_theme"
 # documentation.
 
 html_theme_options = {
-    "github_url": "https://github.com/arviz-devs/arviz",
-    "twitter_url": "https://twitter.com/arviz_devs",
+    "icon_links": [
+        {
+            "name": "GitHub",
+            "url": "https://github.com/arviz-devs/arviz",
+            "icon": "fab fa-github-square",
+        },
+        {
+            "name": "Twitter",
+            "url": "https://twitter.com/arviz_devs",
+            "icon": "fab fa-twitter-square",
+        },
+    ],
+    "navbar_start": ["navbar-logo", "navbar-version"],
+    "use_edit_page_button": False,  # TODO: see how to skip of fix for generated pages
+}
+html_context = {
+    "github_user": "arviz-devs",
+    "github_repo": "arviz",
+    "github_version": "main",
+    "doc_path": "doc/source/",
 }
 
 # Add any paths that contain custom static files (such as style sheets) here,
diff --git a/doc/source/user_guide/Numba.ipynb b/doc/source/user_guide/Numba.ipynb
index f4aa3ac..bdcb456 100644
--- a/doc/source/user_guide/Numba.ipynb
+++ b/doc/source/user_guide/Numba.ipynb
@@ -33,7 +33,6 @@
     "import timeit\n",
     "\n",
     "from arviz.utils import conditional_jit, Numba\n",
-    "from arviz.stats import geweke\n",
     "from arviz.stats.diagnostics import ks_summary"
    ]
   },
@@ -288,9 +287,9 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.6.10"
+   "version": "3.9.2"
   }
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
\ No newline at end of file
+}


Test patch:
diff --git a/arviz/tests/base_tests/test_plots_bokeh.py b/arviz/tests/base_tests/test_plots_bokeh.py
index 396c804..41279a0 100644
--- a/arviz/tests/base_tests/test_plots_bokeh.py
+++ b/arviz/tests/base_tests/test_plots_bokeh.py
@@ -718,9 +718,9 @@ def test_plot_loo_pit_label(models, args):
     if args.get("y_hat") == "str":
         y_hat = "y"
     elif args.get("y_hat") == "DataArray":
-        y_hat = models.model_1.posterior_predictive.y.stack(sample=("chain", "draw"))
+        y_hat = models.model_1.posterior_predictive.y.stack(__sample__=("chain", "draw"))
     elif args.get("y_hat") == "ndarray":
-        y_hat = models.model_1.posterior_predictive.y.stack(sample=("chain", "draw")).values
+        y_hat = models.model_1.posterior_predictive.y.stack(__sample__=("chain", "draw")).values
     else:
         y_hat = None
 
diff --git a/arviz/tests/base_tests/test_stats.py b/arviz/tests/base_tests/test_stats.py
index 4f737d5..260fe30 100644
--- a/arviz/tests/base_tests/test_stats.py
+++ b/arviz/tests/base_tests/test_stats.py
@@ -486,7 +486,7 @@ def test_loo_print(centered_eight, scale):
 def test_psislw(centered_eight):
     pareto_k = loo(centered_eight, pointwise=True, reff=0.7)["pareto_k"]
     log_likelihood = get_log_likelihood(centered_eight)
-    log_likelihood = log_likelihood.stack(sample=("chain", "draw"))
+    log_likelihood = log_likelihood.stack(__sample__=("chain", "draw"))
     assert_allclose(pareto_k, psislw(-log_likelihood, 0.7)[1])
 
 
@@ -544,9 +544,9 @@ def test_loo_pit(centered_eight, args):
     y_hat = args.get("y_hat", None)
     log_weights = args.get("log_weights", None)
     y_arr = centered_eight.observed_data.obs
-    y_hat_arr = centered_eight.posterior_predictive.obs.stack(sample=("chain", "draw"))
-    log_like = get_log_likelihood(centered_eight).stack(sample=("chain", "draw"))
-    n_samples = len(log_like.sample)
+    y_hat_arr = centered_eight.posterior_predictive.obs.stack(__sample__=("chain", "draw"))
+    log_like = get_log_likelihood(centered_eight).stack(__sample__=("chain", "draw"))
+    n_samples = len(log_like.__sample__)
     ess_p = ess(centered_eight.posterior, method="mean")
     reff = np.hstack([ess_p[v].values.flatten() for v in ess_p.data_vars]).mean() / n_samples
     log_weights_arr = psislw(-log_like, reff=reff)[0]
@@ -584,9 +584,9 @@ def test_loo_pit_multidim(multidim_models, args):
     log_weights = args.get("log_weights", None)
     idata = multidim_models.model_1
     y_arr = idata.observed_data.y
-    y_hat_arr = idata.posterior_predictive.y.stack(sample=("chain", "draw"))
-    log_like = get_log_likelihood(idata).stack(sample=("chain", "draw"))
-    n_samples = len(log_like.sample)
+    y_hat_arr = idata.posterior_predictive.y.stack(__sample__=("chain", "draw"))
+    log_like = get_log_likelihood(idata).stack(__sample__=("chain", "draw"))
+    n_samples = len(log_like.__sample__)
     ess_p = ess(idata.posterior, method="mean")
     reff = np.hstack([ess_p[v].values.flatten() for v in ess_p.data_vars]).mean() / n_samples
     log_weights_arr = psislw(-log_like, reff=reff)[0]