code: def plot_head_matrix(z, head_layer, z_max, z_min):
      data = [
          dict(
              type="heatmap",
              z=z,
              colorscale="RdBu",
              reversescale=False,
              zmin=z_min,
              zmax=z_max,
          ),
          dict(
              type="scatter",
              x=[int(head_layer["head_no"])],
              y=[int(head_layer["layer_no"])],
              marker=dict(
                  color="black",
                  symbol="x",
                  size=17,
                  opacity=0.4,
              ),
          ),
      ]
      num_layers, num_heads = np.shape(z)
      layout_ind = copy.deepcopy(layout)
      layout_ind.update({"yaxis": dict(autorange="reversed")})
      layout_ind["yaxis"].update(
          {
              "title": "layer_no",
              "tickmode": "array",
              "ticktext": list(range(1, num_layers + 1)),
              "tickvals": list(range(num_layers)),
          }
      )
      layout_ind["xaxis"] = {
          "title": "head_no",
          "tickmode": "array",
          "ticktext": list(range(1, num_heads + 1)),
          "tickvals": list(range(num_heads)),
      }
      layout_ind.update({"margin": {"t": 20}})
  
      figure = dict(data=data, layout=layout_ind)
  
      return figure
  
file_path: nlp-architect/solutions/InterpreT/application/plotFunc.py
imports: import base64
  import copy
  from io import BytesIO
  import numpy
  import matplotlib.pyplot
  import plotly.graph_objects
label_desc: Use image node_id 1 for calling a global function with example usage: plot_head_matrix(z, head_layer, z_max, z_min) and returns: figure