code: def display_gradients_of_trainable_weights(
      gradients, directory=".", save=False
  ):
      """
      Plot in_channels by out_channels grid of grad heatmaps each of dimensions (filter_h, filter_w)
      :param gradients: dict mapping layers to corresponding gradients (filter_h, filter_w, in_channels, out_channels)
      :param save: bool- if the plot should be saved
      :param directory: string - where to store the activations (if save is True)
      :return: None
      """
      import matplotlib.pyplot as plt
  
      index = 0
      for layer_name, grads in gradients.items():
          if len(grads.shape) != 4:
              print(
                  layer_name,
                  ": Expected dimensions (filter_h, filter_w, in_channels, out_channels). Got ",
                  grads.shape,
              )
              continue
          print(layer_name, grads.shape)
          nrows = grads.shape[-1]
          ncols = grads.shape[-2]
          fig, axes = plt.subplots(nrows, ncols, figsize=(12, 12))
          fig.suptitle(layer_name)
          hmap = None
          for i in range(nrows):
              for j in range(ncols):
                  g = grads[:, :, j, i]
                  hmap = axes[i, j].imshow(
                      g, aspect="auto"
                  )  # May cause distortion in case of in_out channel difference
                  axes[i, j].axis("off")
          fig.subplots_adjust(right=0.8, wspace=0.02, hspace=0.3)
          cbar = fig.add_axes([0.85, 0.15, 0.03, 0.7])
          fig.colorbar(hmap, cax=cbar)
          if save:
              if not os.path.exists(directory):
                  os.makedirs(directory)
              output_filename = os.path.join(
                  directory,
                  "{0}_{1}.png".format(index, layer_name.split("/")[0]),
              )
              plt.savefig(output_filename, bbox_inches="tight")
          else:
              plt.show()
          index += 1
          plt.close(fig)
  
imports: import json
  import os
  from collections import OrderedDict
  import numpy
  import tensorflow
  import tensorflow.keras.backend
  from tensorflow.keras import Sequential
  from tensorflow.keras.models import Model
label_desc: Use image node_id 12 for calling a global function with example usage: display_gradients_of_trainable_weights(gradients, directory, save) without return types