class_docstr: Helper class to create graphs to display memory
class_docstr_tok: ["Helper","class","to","create","graphs","to","display","memory"]
code: def module_parameters(
      self, ax, job_name: str, memory_traces: List[LayerMemoryTrace]
  ):
      module_parameters = [t.module_params for t in memory_traces]
      x, y_forward, y_backward = self._split_forward_backward(
          memory_traces, module_parameters
      )
      ax.plot(x, y_forward, x, y_backward, label=job_name)
      self._y_axis_in_gigabytes(ax)
  
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 7 for calling the MemoryGraphCreator obj's underlying member method code with example usage: obj.module_parameters(ax, job_name, memory_traces) without return types