code: def plot_dependency_matrix(deps):
      """
      Returns a matplotlib figure of a dependency matrix generated by
      `infer_dependency_matrix`.
  
      At a given point, a red square indicates that a given output frame (y-axis)
      was to depend on a given input frame (x-axis).
  
      For example, a fully red image means that all output frames were dependent
      on all the history. This could be the case of a bidirectional RNN, or a
      transformer model, for example.
  
      Arguments
      ---------
      deps : torch.BoolTensor
          Matrix returned by `infer_dependency_matrix` or one in a compatible
          format.
      """
      import matplotlib.pyplot as plt
      from matplotlib.colors import ListedColormap
  
      cmap = ListedColormap(["white", "red"])
  
      fig, ax = plt.subplots()
  
      ax.pcolormesh(
          torch.permute(deps, (1, 0)),
          cmap=cmap,
          vmin=False,
          vmax=True,
          edgecolors="gray",
          linewidth=0.5,
      )
      ax.set_title("Dependency plot")
      ax.set_xlabel("in")
      ax.set_ylabel("out")
      ax.set_aspect("equal")
      return fig
  
imports: import math
  import torch
  from typing import Callable, List
label_desc: Use image node_id 4 for calling a global function with example usage: plot_dependency_matrix(deps) and returns: fig