class_docstr: Helper class to create graphs to display memory
class_docstr_tok: ["Helper","class","to","create","graphs","to","display","memory"]
code: def all_gathered_memory(
      self, ax, job_name: str, memory_traces: List[LayerMemoryTrace]
  ):
      # Plot the all_gathered and cumulative all_gathered memory
      gathered_memory = [t.all_gathered for t in memory_traces]
      cumul_gathered_memory = [
          t.cumul_all_gathered for t in memory_traces
      ]
      x, y_forward, y_backward = self._split_forward_backward(
          memory_traces, gathered_memory
      )
      ax.plot(x, y_forward, x, y_backward, label=job_name)
      ax.plot(x, cumul_gathered_memory, label=job_name)
      self._y_axis_in_gigabytes(ax)
  
      # Adding the name of the layer with max cumulative all_gathered memory
      max_index = np.argmax(cumul_gathered_memory)
      max_trace = memory_traces[max_index]
      max_module = ".".join(
          [
              n
              for n in max_trace.module_name.split(".")
              if not n.startswith("_")
          ]
      )
      ax.set_ylim([None, max_trace.cumul_all_gathered * 1.1])
      x_text, y_text = (
          max(0, max_index * 0.8),
          max_trace.cumul_all_gathered * 1.04,
      )
      ax.text(x_text, y_text, f"{max_module} (fwd)", fontdict=self.font)
  
imports: from contextlib import contextmanager
  from dataclasses import dataclass
  from enum import auto, Enum
  from functools import lru_cache
  from typing import Dict, List, NamedTuple, Tuple, Union
  import numpy
  import torch
  import torch.nn
  from fairscale.nn import FullyShardedDataParallel
  from vissl.utils.visualize import matplotlib_figure_to_image
label_desc: Use image node_id 6 for calling the MemoryGraphCreator obj's underlying member method code with example usage: obj.all_gathered_memory(ax, job_name, memory_traces) without return types