code: def _plot_loss_fn(losses, ax=None, log_scale=True) -> plt.Figure:
      """Plot losses from optimization."""
      if ax is None:
          fig, ax = plt.subplots(
              figsize=(16, 4), constrained_layout=True
          )
      else:
          fig = ax.figure
      flat_losses = losses.reshape((-1, losses.shape[-1])).T
      step = jnp.arange(flat_losses.shape[0])
      ax.plot(step, flat_losses, "-", alpha=0.5)
  
      x_val = int(step.max() * 0.75)
  
      xlim = (x_val, step.max())
      ylim = (
          0.95 * flat_losses.min(),
          1.05 * flat_losses[x_val:].max(),
      )
      axins = ax.inset_axes(
          [0.5, 0.5, 0.47, 0.4],  # This is in axis units, not data
          xlim=xlim,
          ylim=ylim,
          xticklabels=[],
      )
      axins.plot(step[x_val:], flat_losses[x_val:], "-", alpha=0.8)
      if log_scale:
          axins.set_yscale("log")
      ax.indicate_inset_zoom(axins, edgecolor="black")
      if log_scale:
          ax.set_yscale("log")
      ax.set_title("Loss")
      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 6 for calling a global function with example usage: _plot_loss_fn(losses, ax, log_scale) and returns: fig