code: def display_imgs(x, title=None, fignum=None):
      """Display images as a grid."""
      import matplotlib.pyplot as plt  # pylint: disable=import-outside-toplevel,g-import-not-at-top
  
      if not tf.executing_eagerly():
          raise NotImplementedError(
              "`display_imgs` can only be executed eagerly."
          )
  
      def _preprocess(z):
          return np.array(getattr(z, "numpy", lambda: z)())
  
      x = _preprocess(x)
      if title is not None:
          title = _preprocess(title)
      x = np.reshape(x, (-1,) + x.shape[-4:])
      nrows, ncols, h, w, c = x.shape
      x = np.reshape(
          np.transpose(x, [0, 2, 1, 3, 4]), [nrows * h, ncols * w, c]
      )
      plt.ioff()
      subplots_kwargs = dict(
          nrows=1,
          ncols=1,
          figsize=(ncols, max(2, nrows)),
          num=fignum,
          clear=True,
      )
      try:
          fig, axs = plt.subplots(**subplots_kwargs)
      except TypeError:
          subplots_kwargs.pop("clear")
          fig, axs = plt.subplots(**subplots_kwargs)
      axs.imshow(x.squeeze(), interpolation="none", cmap="gray")
      axs.axis("off")
      if title is not None:
          axs.set_title(str(title))
      fig.tight_layout()
      plt.show()
      plt.ion()
      return fig, axs
  
imports: import collections
  import contextlib
  import functools
  import sys
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
label_desc: Use image node_id 1 for calling a global function with example usage: display_imgs(x, title, fignum) and returns: fig,  axs, np