Repository: marimo-team/marimo
Pull request: `Arviz` plotting library fix
## 📝 Summary
This PR addresses the issue with [ArviZ](https://python.arviz.org/en/stable/index.html) plots not displaying correctly in the Marimo output. It implements a custom formatter for ArviZ objects, specifically handling numpy arrays containing matplotlib `Axes` objects along with `az.InferenceData`. 

Fixes #1033

## 🔍 Description of Changes
- Implemented a new `ArviZFormatter`.
- Added methods to detect and handle numpy arrays containing matplotlib Axes objects.
- Implemented formatting methods for various ArviZ objects, including InferenceData, numpy arrays and matplotlib Figures.
- Created utility methods to extract axes information and convert plots to HTML format.
- Ensures consistency in plot rendering: This includes adapting plots like `arviz.plot_bf`, `arviz.plot_density`, and others to correctly display without the need for `plt.show()` and handling `matplotlib.axes.Axes` objects properly.
 - Improved handling of InferenceData across various plotting functions, ensuring that plots render as expected in both Matplotlib and Bokeh, following Jupyter's output behavior for better compatibility.

Before this fix, most plots in ArviZ were experiencing rendering issues due to improper handling of the backend and the necessity of explicityly calling `plt.show()` or `plt.gca()`. This fix ensures that the output of plotting functions behaves similarly to Jupyter notebooks.

## 📋 Checklist
- [x] I have read the [contributor guidelines](https://github.com/marimo-team/marimo/blob/main/CONTRIBUTING.md).
- [ ] For large changes, or changes that affect the public API: this change was discussed or approved through an issue, on [Discord](https://discord.gg/JE7nhX6mD8), or the community [discussions](https://github.com/marimo-team/marimo/discussions) (Please provide a link if applicable).
- [x] I have added tests for the changes made.
- [x] I have run the code and verified that it works as expected.

## 📜 Reviewers
@akshayka OR @mscolnick


Problem statement:
Arviz plots not being displayed
### Describe the bug

I am using the arviz library with Pymc and the plots are not being displayed. All I see is the axis infos but not the plots.

### Environment

`
{
  "marimo": "0.3.7",
  "OS": "Darwin",
  "OS Version": "23.3.0",
  "Processor": "i386",
  "Python Version": "3.8.3",
  "Binaries": {
    "Browser": "123.0.6312.87",
    "Node": "--"
  },
  "Requirements": {
    "click": "8.1.7",
    "importlib-resources": "6.1.1",
    "jedi": "0.19.1",
    "markdown": "3.6",
    "pymdown-extensions": "10.7.1",
    "pygments": "2.17.2",
    "tomlkit": "0.12.4",
    "uvicorn": "0.29.0",
    "starlette": "0.37.2",
    "websocket": "missing",
    "typing-extensions": "4.8.0",
    "black": "24.3.0"
  }
}`

### Code to reproduce

```python
import marimo

__generated_with = "0.3.7"
app = marimo.App(width="medium")


@app.cell
def __():
    import marimo as mo
    return mo,


@app.cell
def __():
    import pymc as pm
    import numpy as np
    import pandas as pd
    return np, pd, pm


@app.cell
def __(np):
    seed = 134
    rng = np.random.default_rng(seed=seed)
    return rng, seed


@app.cell
def __(mo):
    mo.md(
        "We create a series of kmers that could be used a preferential cutting sites"
    )
    return


@app.cell
def __():
    from random import choices

    n = 0
    N = 4**4

    kmers = set()

    while n < N:
        kmer = "".join(choices(k=4, population="ATCG"))
        kmers.add(kmer)
        n = len(kmers)
    return N, choices, kmer, kmers, n


@app.cell
async def __(choices, kmers, np, rng):
    k = 7
    subkmers = choices(population=list(kmers), k=k)

    frac = np.array([0.2, 0.1, 0.1, 0.1, 0.4, 0.05, 0.05])
    if np.sum(frac) == 1 and k == len(frac):
        print("OK")
    else:
        print("Error")

    true_c = 3

    # number of reads
    readsLen = 10
    reads = np.arange(readsLen)

    true_p = rng.dirichlet(size=readsLen, alpha=frac * true_c)

    ## get the counts
    totalcount = 34
    obs_counts = np.vstack(
        [rng.multinomial(n=totalcount, pvals=pi) for pi in true_p]
    )
    obs_counts
    return (
        frac,
        k,
        obs_counts,
        reads,
        readsLen,
        subkmers,
        totalcount,
        true_c,
        true_p,
    )


@app.cell
def __(obs_counts, true_p):
    print(
        f"shape of true_p {true_p.shape} \nshape of obs_count {obs_counts.shape}"
    )
    obs_counts.shape
    return


@app.cell
def __(k, np, obs_counts, pm, reads, subkmers, totalcount):
    coords = {"reads": reads, "subkmers": subkmers}

    with pm.Model(coords=coords) as multidirimarg:
        fraction = pm.Dirichlet("frac", a=np.ones(k), dims="subkmers")
        conc = pm.Lognormal("conc", mu=1, sigma=1)
        counts = pm.DirichletMultinomial(
            "counts",
            n=totalcount,
            a=fraction * conc,
            observed=obs_counts,
            dims=("reads", "subkmers"),
        )

    pm.model_to_graphviz(multidirimarg)
    with multidirimarg:
        trace_dm_marginalized = pm.sample(chains=4)
    return (
        conc,
        coords,
        counts,
        fraction,
        multidirimarg,
        trace_dm_marginalized,
    )


@app.cell
def __(trace_dm_marginalized):
    import arviz as az

    trace_plot = az.plot_trace(
        data=trace_dm_marginalized, var_names=["frac", "conc"]
    )
    trace_plot
    return az, trace_plot


@app.cell
def __():
    import matplotlib.pyplot as plt

    plt.plot([1, 2])
    # plt.gca() gets the current `Axes`
    plt.gca()
    return plt,


if __name__ == "__main__":
    app.run()
```

Hints from issue discussion:
Adding support for arviz would involve adding a formatter, here: https://github.com/marimo-team/marimo/tree/main/marimo/_output/formatters

This is a good first issue. For anyone interested, you can look at the other formatters for examples.
Working on this. Will try to resolve and refer to [formatters](https://github.com/marimo-team/marimo/tree/main/marimo/_output/formatters) provided above.
awesome, thanks @Haleshot!
Trying out different variations; saw how seaborn uses it's plotting functions; plotly and leafmap too. 
@lborcard Current using your code provided above to understand the return types thrown; assigned a variable to the `plot_trace` method and it seems it is of the `numpy.ndarray` type. 

Will try different approaches to see how I can double down on a solution,
After extensive investigation through the library's docs and other resources, I've discovered a simple fix for the ArviZ formatter issue. It appears that adding `plt.show()` after [any ArviZ plot function](https://python.arviz.org/en/stable/api/plots.html) call resolves the display problem. This approach leverages the fact that ArviZ uses matplotlib as its default backend (unless otherwise specified in parameters through functions to use `Bokeh`), as explained [here](https://python.arviz.org/en/v0.18.0/user_guide/plotting_with_matplotlib.html).

For example, the ArviZ documentation [showcases this usage](https://python.arviz.org/en/stable/examples/plot_trace.html#traceplot) of `plt.show()` in their examples (which is where I came across the simple fix). (_FYI @lborcard_)

The solution is straightforward: append `plt.show()` to the end of the cell block containing the ArviZ plot function (trace_plot in this case - considering the code given above for reproduction).

Given this, I'm considering posting this as a simpler alternative to implementing a complex formatter. The current formatter implementation, which applies changes to all `np.ndarray` formats, is causing unintended effects on other libraries, particularly matplotlib (posted on the discord server as a thread in #contributing).

I'm planning to update my current local changes to reflect this simpler approach.

@akshayka @mscolnick (_Apologize for the ping_)
Do you think it's still worthwhile to create a formatter explicitly for ArviZ which handles the `plt.show()` on itself without the user having to explicitly mention it in their code (cell blocks) to automatically apply this change for ArviZ library usage? Or should we opt for the simpler, more explicit solution of recommending users add `plt.show()` after ArviZ plot calls (that too only for the ones that are of the `np.ndarray type` like plot_trace?

Your thoughts on this approach would be greatly appreciated.

Commit message: [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Patch to review:
diff --git a/marimo/_output/formatters/arviz_formatters.py b/marimo/_output/formatters/arviz_formatters.py
new file mode 100644
index 00000000000..a8050cb950e
--- /dev/null
+++ b/marimo/_output/formatters/arviz_formatters.py
@@ -0,0 +1,151 @@
+from __future__ import annotations
+
+from typing import Any
+
+from marimo._messaging.mimetypes import KnownMimeType
+from marimo._output.formatters.formatter_factory import FormatterFactory
+
+
+class ArviZFormatter(FormatterFactory):
+    import matplotlib.pyplot as plt
+    import numpy as np
+
+    @staticmethod
+    def package_name() -> str:
+        return "arviz"
+
+    def register(self) -> None:
+        import arviz as az
+        import matplotlib.pyplot as plt
+        import numpy as np
+
+        from marimo._output import formatting
+
+        @formatting.formatter(az.InferenceData)
+        def _format_inference_data(
+            data: az.InferenceData,
+        ) -> tuple[KnownMimeType, str]:
+            return ("text/plain", str(data))
+
+        @formatting.formatter(np.ndarray)
+        def _format_ndarray(
+            arr: np.ndarray,
+        ) -> tuple[KnownMimeType, str]:
+            return self.format_numpy_axes(arr)
+
+        @formatting.formatter(dict)
+        def _format_dict(
+            d: dict,
+        ) -> tuple[KnownMimeType, str]:
+            return self.format_dict_with_plot(d)
+
+        @formatting.formatter(plt.Figure)
+        def _format_figure(
+            fig: plt.Figure,
+        ) -> tuple[KnownMimeType, str]:
+            return self.format_figure(fig)
+
+        @formatting.formatter(object)
+        def _format_arviz_plot(
+            obj: Any,
+        ) -> tuple[KnownMimeType, str]:
+            return self.format_arviz_plot(obj)
+
+    @classmethod
+    def format_numpy_axes(cls, arr: np.ndarray) -> tuple[KnownMimeType, str]:
+        import matplotlib.pyplot as plt
+
+        # Check if array contains axes (to render plots) or not
+        if arr.dtype == object and cls._contains_axes(arr):
+            fig = plt.gcf()
+            if fig.get_axes():  # Only process if there are axes to show
+                axes_info = cls._get_axes_info(fig)
+                plot_html = cls._get_plot_html(fig)
+                plt.close(fig)  # Safely close the figure after saving
+                combined_html = f"<pre>{axes_info}</pre><br>{plot_html}"
+                return ("text/html", combined_html)
+        # Fallback to plain text if no axes or plot are present
+        return ("text/plain", str(arr))
+
+    @staticmethod
+    def _contains_axes(arr: np.ndarray) -> bool:
+        import matplotlib.pyplot as plt
+
+        """
+        Check if the numpy array contains any matplotlib Axes objects.
+        To ensure performance for large arrays, we limit the check to the
+        first 100 items. This should be sufficient for most use cases
+        while avoiding excessive computation time.
+        """
+        # Cap the number of items to check for performance reasons
+        MAX_ITEMS_TO_CHECK = 100
+
+        if arr.ndim == 1:
+            # For 1D arrays, check up to MAX_ITEMS_TO_CHECK items
+            return any(
+                isinstance(item, plt.Axes) for item in arr[:MAX_ITEMS_TO_CHECK]
+            )
+        elif arr.ndim == 2:
+            # For 2D arrays, check up to MAX_ITEMS_TO_CHECK items in total
+            items_checked = 0
+            for row in arr:
+                for item in row:
+                    if isinstance(item, plt.Axes):
+                        return True
+                    items_checked += 1
+                    if items_checked >= MAX_ITEMS_TO_CHECK:
+                        return False
+        return False
+
+    @staticmethod
+    def _get_plot_html(fig: plt.Figure) -> str:
+        import base64
+        from io import BytesIO
+
+        buf = BytesIO()
+        fig.savefig(buf, format="png", bbox_inches="tight")  # Retain default
+        data = base64.b64encode(buf.getbuffer()).decode("ascii")
+        return f"<img src='data:image/png;base64,{data}'/>"
+
+    @classmethod
+    def format_dict_with_plot(cls, d: dict) -> tuple[KnownMimeType, str]:
+        import matplotlib.pyplot as plt
+
+        str_repr = str(d)
+        fig = plt.gcf()
+        if fig.get_axes():
+            axes_info = cls._get_axes_info(fig)
+            plot_html = cls._get_plot_html(fig)
+            plt.close(fig)
+            combined_html = (
+                f"<pre>{str_repr}\n{axes_info}</pre><br>" f"{plot_html}"
+            )
+            return ("text/html", combined_html)
+        return ("text/plain", str_repr)
+
+    @classmethod
+    def format_figure(cls, fig: plt.Figure) -> tuple[KnownMimeType, str]:
+        import matplotlib.pyplot as plt
+
+        axes_info = cls._get_axes_info(fig)
+        plot_html = cls._get_plot_html(fig)
+        plt.close(fig)
+        combined_html = f"<pre>{axes_info}</pre><br>{plot_html}"
+        return ("text/html", combined_html)
+
+    @classmethod
+    def format_arviz_plot(cls, result: Any) -> tuple[KnownMimeType, str]:
+        import matplotlib.pyplot as plt
+        import numpy as np
+
+        if isinstance(result, plt.Figure):
+            return cls.format_figure(result)
+        elif isinstance(result, np.ndarray):
+            return cls.format_numpy_axes(result)
+        elif isinstance(result, dict):
+            return cls.format_dict_with_plot(result)
+        else:
+            fig = plt.gcf()
+            if fig.get_axes():
+                return cls.format_figure(fig)
+            return ("text/plain", str(result))
diff --git a/marimo/_output/formatters/formatters.py b/marimo/_output/formatters/formatters.py
index b93018918b9..2aa7aaf92f6 100644
--- a/marimo/_output/formatters/formatters.py
+++ b/marimo/_output/formatters/formatters.py
@@ -7,6 +7,7 @@
 from marimo._config.config import Theme
 from marimo._output.formatters.altair_formatters import AltairFormatter
 from marimo._output.formatters.anywidget_formatters import AnyWidgetFormatter
+from marimo._output.formatters.arviz_formatters import ArviZFormatter
 from marimo._output.formatters.bokeh_formatters import BokehFormatter
 from marimo._output.formatters.cell import CellFormatter
 from marimo._output.formatters.df_formatters import (
@@ -44,6 +45,7 @@
     IPythonFormatter.package_name(): IPythonFormatter(),
     IPyWidgetsFormatter.package_name(): IPyWidgetsFormatter(),
     AnyWidgetFormatter.package_name(): AnyWidgetFormatter(),
+    ArviZFormatter.package_name(): ArviZFormatter(),
     TqdmFormatter.package_name(): TqdmFormatter(),
     LetsPlotFormatter.package_name(): LetsPlotFormatter(),
     SympyFormatter.package_name(): SympyFormatter(),
diff --git a/marimo/_smoke_tests/bugs/1033.py b/marimo/_smoke_tests/bugs/1033.py
new file mode 100644
index 00000000000..62ad3c978bf
--- /dev/null
+++ b/marimo/_smoke_tests/bugs/1033.py
@@ -0,0 +1,334 @@
+# /// script
+# requires-python = ">=3.11"
+# dependencies = [
+#     "arviz",
+#     "marimo",
+#     "xarray",
+#     "matplotlib",
+#     "scipy",
+#     "numpy",
+# ]
+# ///
+
+import marimo
+
+__generated_with = "0.8.15"
+app = marimo.App(width="medium")
+
+
+@app.cell
+def __(az):
+    _data = az.load_arviz_data('centered_eight')
+    az.plot_autocorr(_data)
+    # [az.plot_autocorr(_data), 
+    # type(az.plot_autocorr(_data))]
+    return
+
+
+@app.cell
+def __(az, np):
+    _idata = az.from_dict(posterior={"a":np.random.normal(1, 0.5, 5000)},
+        prior={"a":np.random.normal(0, 1, 5000)})
+    az.plot_bf(_idata, var_name="a", ref_val=0)
+    # [az.plot_bf(_idata, var_name="a", ref_val=0),
+    # type(az.plot_bf(_idata, var_name="a", ref_val=0))]
+    return
+
+
+@app.cell
+def __(az, np):
+    _data = az.load_arviz_data("regression1d")
+    # [az.plot_bpv(_data, kind="t_stat", t_stat=lambda x:np.percentile(x, q=50, axis=-1)), type(az.plot_bpv(_data, kind="t_stat", t_stat=lambda x:np.percentile(x, q=50, axis=-1)))]
+    az.plot_bpv(_data, kind="t_stat", t_stat=lambda x:np.percentile(x, q=50, axis=-1))
+    return
+
+
+@app.cell
+def __(az):
+    _model_compare = az.compare({'Centered 8 schools': az.load_arviz_data('centered_eight'),
+                     'Non-centered 8 schools': az.load_arviz_data('non_centered_eight')})
+    # [az.plot_compare(_model_compare), type(az.plot_compare(_model_compare))]
+    az.plot_compare(_model_compare)
+    return
+
+
+@app.cell
+def __(az):
+    _centered = az.load_arviz_data('centered_eight')
+    _non_centered = az.load_arviz_data('non_centered_eight')
+    # [az.plot_density([_centered, _non_centered]), type(az.plot_density([_centered, _non_centered]))]
+    az.plot_density([_centered, _non_centered])
+    return
+
+
+@app.cell
+def __():
+    # az.clear_data_home()
+    # _data = az.load_arviz_data('rugby')
+    # [az.plot_dist_comparison(_data, var_names=["defs"], coords={"team" : ["Italy"]}), type(az.plot_dist_comparison(_data, var_names=["defs"], coords={"team" : ["Italy"]}))]
+    return
+
+
+@app.cell
+def __(az, np):
+    _values = np.random.normal(0, 1, 500)
+    # [az.plot_dot(_values), type(az.plot_dot(_values))]
+    az.plot_dot(_values)
+    return
+
+
+@app.cell
+def __(az, norm, np):
+    _sample = norm(0,1).rvs(1000)
+    _npoints = 100
+    # [az.plot_ecdf(_sample, eval_points=np.linspace(_sample.min(), _sample.max(), _npoints)), type(az.plot_ecdf(_sample, eval_points=np.linspace(_sample.min(), _sample.max(), _npoints)))]
+    az.plot_ecdf(_sample, eval_points=np.linspace(_sample.min(), _sample.max(), _npoints))
+    return
+
+
+@app.cell
+def __(az):
+    _idata1 = az.load_arviz_data("centered_eight")
+    _idata2 = az.load_arviz_data("non_centered_eight")
+    # [az.plot_elpd(
+    #     {"centered model": _idata1, "non centered model": _idata2},
+    #     xlabels=True
+    # ), type(az.plot_elpd(
+    #     {"centered model": _idata1, "non centered model": _idata2},
+    #     xlabels=True
+    # ))]
+    az.plot_elpd(
+        {"centered model": _idata1, "non centered model": _idata2},
+        xlabels=True
+    )
+    return
+
+
+@app.cell
+def __(az):
+    _data = az.load_arviz_data('centered_eight')
+    # [az.plot_energy(_data), type(az.plot_energy(_data))]
+    az.plot_energy(_data)
+    return
+
+
+@app.cell
+def __(az):
+    _idata = az.load_arviz_data("centered_eight")
+    _coords = {"school": ["Choate", "Lawrenceville"]}
+    az.plot_ess(
+        _idata, kind="local", var_names=["mu", "theta"], coords=_coords
+    )
+    # [az.plot_ess(
+    #     _idata, kind="local", var_names=["mu", "theta"], coords=_coords
+    # ), type(az.plot_ess(
+    #     _idata, kind="local", var_names=["mu", "theta"], coords=_coords
+    # ))]
+    return
+
+
+@app.cell
+def __(az):
+    _non_centered_data = az.load_arviz_data('non_centered_eight')
+    _axes = az.plot_forest(_non_centered_data,
+                               kind='forestplot',
+                               var_names=["^the"],
+                               filter_vars="regex",
+                               combined=True,
+                               figsize=(9, 7))
+    _axes[0].set_title('Estimated theta for 8 schools model')
+    # [_axes, type(_axes)]
+    return
+
+
+@app.cell
+def __(az, np):
+    # time-steps random walk
+    _x_data =np.arange(0,100)
+    # Mean random walk
+    _mu = np.zeros(100)
+    for i in _x_data: _mu[i] = _mu[i-1] + np.random.normal(0, 1, 1)
+    # Simulated pp samples form the random walk time series
+    _y_data = np.random.normal(2 + _mu * 0.5, 0.5, size = (2, 50, 100))
+    # [az.plot_hdi(_x_data, _y_data), type(az.plot_hdi(_x_data, _y_data))]
+    az.plot_hdi(_x_data, _y_data)
+    return i,
+
+
+@app.cell
+def __(az, np):
+    _non_centered = az.load_arviz_data('non_centered_eight')
+    _mu_posterior = np.concatenate(_non_centered.posterior["mu"].values)
+    _tau_posterior = np.concatenate(_non_centered.posterior["tau"].values)
+    # [az.plot_kde(_mu_posterior), type(az.plot_kde(_mu_posterior))]
+    az.plot_kde(_mu_posterior)
+    return
+
+
+@app.cell
+def __(az):
+    _radon = az.load_arviz_data("radon")
+    _loo_radon = az.loo(_radon, pointwise=True)
+    # [az.plot_khat(_loo_radon, show_bins=True), type(az.plot_khat(_loo_radon, show_bins=True))]
+    az.plot_khat(_loo_radon, show_bins=True)
+    return
+
+
+@app.cell
+def __(az):
+    _idata = az.load_arviz_data("radon")
+    # [az.plot_loo_pit(idata=_idata, y="y"), type(az.plot_loo_pit(idata=_idata, y="y"))]
+    az.plot_loo_pit(idata=_idata, y="y")
+    return
+
+
+@app.cell
+def __(az, np, xr):
+    _idata = az.load_arviz_data('regression1d')
+    _x = xr.DataArray(np.linspace(0, 1, 100))
+    _idata.posterior["y_model"] = _idata.posterior["intercept"] + _idata.posterior["slope"]*_x
+
+    # [az.plot_lm(idata=_idata, y="y", x=_x), type(az.plot_lm(idata=_idata, y="y", x=_x))]
+    az.plot_lm(idata=_idata, y="y", x=_x)
+    return
+
+
+@app.cell
+def __(az):
+    _idata = az.load_arviz_data("centered_eight")
+    _coords = {"school": ["Deerfield", "Lawrenceville"]}
+    # [az.plot_mcse(
+    #     _idata, var_names=["mu", "theta"], coords=_coords
+    # ), type(az.plot_mcse(
+    #     _idata, var_names=["mu", "theta"], coords=_coords
+    # ))]
+    az.plot_mcse(
+        _idata, var_names=["mu", "theta"], coords=_coords
+    )
+    return
+
+
+@app.cell
+def __(az):
+    _centered = az.load_arviz_data('centered_eight')
+    _coords = {'school': ['Choate', 'Deerfield']}
+    # [az.plot_pair(_centered,
+    #             var_names=['theta', 'mu', 'tau'],
+    #             kind='kde',
+    #             coords=_coords,
+    #             divergences=True,
+    #             textsize=18), type(az.plot_pair(_centered,
+    #             var_names=['theta', 'mu', 'tau'],
+    #             kind='kde',
+    #             coords=_coords,
+    #             divergences=True,
+    #             textsize=18))]
+    az.plot_pair(_centered,
+                var_names=['theta', 'mu', 'tau'],
+                kind='kde',
+                coords=_coords,
+                divergences=True,
+                textsize=18)
+    return
+
+
+@app.cell
+def __(az, plt):
+    _data = az.load_arviz_data('centered_eight')
+    # [az.plot_parallel(_data, var_names=["mu", "tau"]), type(az.plot_parallel(_data, var_names=["mu", "tau"]))]
+    az.plot_parallel(_data, var_names=["mu", "tau"])
+    plt.show()
+    return
+
+
+@app.cell
+def __(az):
+    _data = az.load_arviz_data('centered_eight')
+    # [az.plot_posterior(_data), type(az.plot_posterior(_data))]
+    az.plot_posterior(_data)
+    return
+
+
+@app.cell
+def __(az):
+    _data = az.load_arviz_data('radon')
+    # [az.plot_ppc(_data, data_pairs={"y":"y"}), type(az.plot_ppc(_data, data_pairs={"y":"y"}))]
+    az.plot_ppc(_data, data_pairs={"y":"y"})
+    # plt.gca()
+    return
+
+
+@app.cell
+def __(az, plt):
+    _data = az.load_arviz_data('centered_eight')
+    # [az.plot_rank(_data), type(az.plot_rank(_data))]
+    az.plot_rank(_data)
+    plt.gca()
+    return
+
+
+@app.cell
+def __(az):
+    _idata = az.load_arviz_data('classification10d')
+    # [az.plot_separation(idata=_idata, y='outcome', y_hat='outcome', figsize=(8, 1)), type(az.plot_separation(idata=_idata, y='outcome', y_hat='outcome', figsize=(8, 1)))]
+    az.plot_separation(idata=_idata, y='outcome', y_hat='outcome', figsize=(8, 1))
+    return
+
+
+@app.cell
+def __(az):
+    _data = az.load_arviz_data('non_centered_eight')
+    _coords = {'school': ['Choate', 'Lawrenceville']}
+    # [az.plot_trace(_data, var_names=('theta'), filter_vars="like", coords=_coords), type(az.plot_trace(_data, var_names=('theta'), filter_vars="like", coords=_coords))]
+    az.plot_trace(_data, var_names=('theta'), filter_vars="like", coords=_coords)
+    return
+
+
+@app.cell
+def __(az, np):
+    _nchains, _ndraws = (4, 500)
+    _obs_data = {
+        "y": 2 * np.arange(1, 9) + 3,
+        "z": 2 * np.arange(8, 12) + 3,
+    }
+    _posterior_predictive = {
+        "y": np.random.normal(
+            (_obs_data["y"] * 1.2) - 3, size=(_nchains, _ndraws, len(_obs_data["y"]))
+        ),
+        "z": np.random.normal(
+            (_obs_data["z"] * 1.2) - 3, size=(_nchains, _ndraws, len(_obs_data["z"]))
+        ),
+     }
+    _idata = az.from_dict(
+        observed_data=_obs_data,
+        posterior_predictive=_posterior_predictive,
+        _coords={"obs_dim": np.arange(1, 9), "pred_dim": np.arange(8, 12)},
+        dims={"y": ["obs_dim"], "z": ["pred_dim"]},
+    )
+    ax = az.plot_ts(idata=_idata, y="y", y_holdout="z")
+    # [ax, type(ax)]
+    ax
+    return ax,
+
+
+@app.cell
+def __(az):
+    _data = az.load_arviz_data('centered_eight')
+    az.plot_violin(_data)
+    return
+
+
+@app.cell(hide_code=True)
+def __():
+    # import libraries
+    import marimo as mo
+    import numpy as np
+    import arviz as az
+    import matplotlib.pyplot as plt
+    import xarray as xr
+    from scipy.stats import uniform, norm
+    return az, mo, norm, np, plt, uniform, xr
+
+
+if __name__ == "__main__":
+    app.run()


Code review comments:
diff_hunk: @@ -0,0 +1,151 @@
+from __future__ import annotations
+
+from typing import Any
+
+from marimo._messaging.mimetypes import KnownMimeType
+from marimo._output.formatters.formatter_factory import FormatterFactory
+
+
+class ArviZFormatter(FormatterFactory):
+    import matplotlib.pyplot as plt
+    import numpy as np
path: marimo/_output/formatters/arviz_formatters.py
text: @user1:
These should be removed