def calculate_memory(graph, parameter=True, saved=True, peak_activation=False):
  """Calculate memory consumption of the graph based on flags.

  Args:
    graph(nx.DiGraph): Input graph with attributes for memory usage.
    parameter(bool): If True, only count parameter memory.
    saved(bool): If True, only count activation memory.
    peak_activation(bool): If True, only count peak activation memory.

  Returns: Memory consumed by the graph.

  """
  total_mem = 0
  peak_activation_mem = 0
  parameter_mem = 0

  for node in graph.nodes:
    total_mem += graph.nodes[node]['saved_mem']
    total_mem += graph.nodes[node]['parameter_mem']

    peak_activation_mem = max(peak_activation_mem,
                              graph.nodes[node]['saved_mem'])

    parameter_mem += graph.nodes[node]['parameter_mem']

  total = parameter and saved
  if peak_activation:
    return peak_activation_mem
  if total:
    return total_mem
  if parameter:
    return parameter_mem
  if saved:
    return total_mem - parameter_mem