code: def plot_results(
      dates_preds: Union[Sequence[np.datetime64], jax.Array],
      preds: jax.Array,
      *,
      dates_test: Union[
          Sequence[np.datetime64], jax.Array, None
      ] = None,
      y_test: Optional[jax.Array] = None,
      p2_5: Optional[jax.Array] = None,
      p50: Optional[jax.Array] = None,
      p97_5: Optional[jax.Array] = None,
      dates_train: Union[
          Sequence[np.datetime64], jax.Array, None
      ] = None,
      y_train: Optional[jax.Array] = None,
      diagnostics: Optional[Dict[str, jax.Array]] = None,
      log_scale: bool = False,
      left_limit: int = 24 * 7 * 2,
      right_limit: int = 24 * 7 * 2,
  ) -> plt.Figure:
      """Plot the results of `fit_bnn_map`."""
      if diagnostics is None:
          diagnostics = {}
      if diagnostics.get("loss") is not None:
          fig, (aux_ax, res_ax) = plt.subplots(
              figsize=(16, 6), nrows=2, constrained_layout=True
          )
          _plot_loss_fn(
              diagnostics["loss"], ax=aux_ax, log_scale=log_scale
          )
      elif diagnostics.get("noise_scale") is not None:
          fig = plt.figure(figsize=(16, 6), constrained_layout=True)
          axes = fig.subplot_mosaic("tttth;rrrrr")
          _plot_noise_fn(
              diagnostics["noise_scale"], axes["t"], axes["h"]
          )
          res_ax = axes["r"]
      else:
          fig, res_ax = plt.subplots(
              figsize=(16, 3), constrained_layout=True
          )
  
      for idx, p in enumerate(preds):
          res_ax.plot(
              dates_preds,
              p,
              "k-",
              alpha=0.1,
              label="Particle predictions" if idx == 0 else None,
          )
  
      color = "steelblue"
      if p50 is not None:
          res_ax.plot(
              dates_preds,
              p50,
              "-",
              lw=5,
              color=color,
              label="Prediction",
          )
      if p97_5 is not None and p2_5 is not None:
          res_ax.plot(
              dates_preds,
              p97_5,
              "-",
              lw=3,
              color=color,
              label="Upper/lower bound",
          )
          res_ax.plot(dates_preds, p2_5, "-", lw=3, color=color)
          res_ax.fill_between(
              dates_preds, p2_5, p97_5, color=color, alpha=0.2
          )
  
      data_kwargs = {"ms": 7, "mec": "k", "mew": 2}
      if dates_train is not None and y_train is not None:
          res_ax.plot(
              dates_train,
              y_train,
              "o",
              mfc="red",
              label="Train data",
              **data_kwargs,
          )
      if dates_test is not None and y_test is not None:
          res_ax.plot(
              dates_test,
              y_test,
              "o",
              mfc="green",
              label="Test data",
              **data_kwargs,
          )
      res_ax.set_title("Predictions")
      res_ax.legend()
      left_limit = min(len(dates_preds) - len(dates_test), left_limit)
      right_limit = min(len(dates_test) - 1, right_limit)
      first_test_point = np.where(dates_preds == dates_test[0])[0][0]
      # TODO(ursk): Rather than modifying xlim, don't plot invisible points at all.
      res_ax.set_xlim(
          [
              dates_preds[first_test_point - left_limit],
              dates_preds[first_test_point + right_limit],
          ]
      )
      return fig
  
imports: import functools
  from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
  import bayeux
  import jax
  import jax.numpy
  from jaxtyping import PyTree
  import matplotlib.pyplot
  import numpy
  import pandas
  from tensorflow_probability.python.experimental.autobnn import bnn
  from tensorflow_probability.python.experimental.autobnn import util
  from tensorflow_probability.python.experimental.timeseries import metrics
  import tensorflow_probability.substrates.jax
label_desc: Use image node_id 10 for calling a global function with example usage: plot_results(dates_preds, preds) and returns: fig