text: # Utilities to figure out memory usage of run call
  #
  # Usage:
  #   import mem_util
  #   run_metadata = tf.RunMetadata()
  #   options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
  #   sess.run(tensor, options=options, run_metadata=run_metadata)
  #   print(mem_util.peak_memory(run_metadata))
  #
  #  To print memory usage for particular device:
  #   print(mem_util.peak_memory(run_metadata)["/gpu:0"])
  
  
  # Developer notes:
  # RunMetadata
  # https://github.com/tensorflow/tensorflow/blob/fc49f43817e363e50df3ff2fd7a4870ace13ea13/tensorflow/core/protobuf/config.proto#L349
  #
  # StepStats (run_metadata.step_stats)
  # https://github.com/tensorflow/tensorflow/blob/a2d9b3bf5f9e96bf459074d079b01e1c74b25afa/tensorflow/core/framework/step_stats.proto
  #
  # NodeExecStats (run_metadata.step_stats.dev_stats[0].step_stats[0])
  # https://github.com/tensorflow/tensorflow/blob/a2d9b3bf5f9e96bf459074d079b01e1c74b25afa/tensorflow/core/framework/step_stats.proto#L52
  
  
  # Note, there are several methods of tracking memory allocation. There's
  # requested bytes, and allocated bytes. Allocator may choose to give more bytes
  # than is requested. Currently allocated_bytes is used to give more realistic
  # results
  #
  # allocation_description {
  #   requested_bytes: 1000000
  #   allocated_bytes: 1000192
  #   allocator_name: "GPU_0_bfc"
  #
  #
  # There's also additional field in NodeExecStats which tracks allocator state
  # node_stats {
  #   node_name: "Bn_1/value"
  #   all_start_micros: 1512081861177033
  #   op_start_rel_micros: 1
  #   op_end_rel_micros: 3
  #   all_end_rel_micros: 5
  #   memory {
  #     allocator_name: "GPU_0_bfc"
  #     allocator_bytes_in_use: 3072
  #   }
  #
  # Additionally one could use LOG_MEMORY messages to get memory allocation info
  # See multiple_memory_obtain_test.py for details on using these additional
  # methods
  
  
  def peak_memory(run_metadata):
      """Return dictionary of peak memory usage (bytes) for each device.
  
    {"cpu:0": 20441, ...
    """
  
      assert run_metadata != None
      assert hasattr(run_metadata, "step_stats")
      assert hasattr(run_metadata.step_stats, "dev_stats")
  
      dev_stats = run_metadata.step_stats.dev_stats
      result = {}
      for dev_stat in dev_stats:
          device_name = _simplify_device_name(dev_stat.device)
          result[device_name] = _peak_from_nodestats(dev_stat.node_stats)
      return result
  
  
  def _timeline_from_nodestats(nodestats):
      """Return sorted memory allocation records from list of nodestats
    [NodeExecStats, NodeExecStats...], it's the
    run_metadata.step_stats.dev_stats[0].step_stats object.
  
    Timeline looks like this:
  
   timestamp         nodename  mem delta, allocator name
   [1509481813012895, 'concat', 1000496, 'cpu'],
   [1509481813012961, 'a04', -1000000, 'cpu'],
   [1509481813012968, 'TanhGrad', 0, 'cpu'], 
  
    0 memory allocation is reported for nodes without allocation_records
    """
  
      lines = []
      if not nodestats:
          return []
      for node in nodestats:
          for mem in node.memory:  # can have both cpu and gpu allocator for op
              try:
                  records = mem.allocation_records
              except:
                  records = []
              allocator = mem.allocator_name
              if len(records) > 0:
                  for record in records:
                      line = [
                          record.alloc_micros,
                          node.node_name,
                          record.alloc_bytes,
                          allocator,
                      ]
                      lines.append(line)
              else:
                  output_bytes = -1
                  try:
                      output_bytes = node.output[
                          0
                      ].tensor_description.allocation_description.requested_bytes
                  except:
                      pass
                  line = [node.all_start_micros, node.node_name, 0, "unknown"]
                  lines.append(line)
  
      def first_key(x):
          return x[0]
  
      return sorted(lines, key=first_key)
  
  
  # todo: get rid of "timeline_from_nodestats"
  
  
  def _position_of_largest(my_list):
      """Return index of largest entry """
      import operator
  
      index, value = max(enumerate(my_list), key=operator.itemgetter(1))
      return index
  
  
  def _peak_from_nodestats(nodestats):
      """Given a list of NodeExecStats messages, construct memory timeline."""
  
      timeline = _timeline_from_nodestats(nodestats)
      timestamps = []
      data = []
  
      total_memory = 0
      peak_memory = total_memory
      for record in timeline:
          timestamp, kernel_name, allocated_bytes, allocator_type = record
          allocated_bytes = int(allocated_bytes)
          total_memory += allocated_bytes
          peak_memory = max(total_memory, peak_memory)
      return peak_memory
  
  
  def _print_parsed_timeline(timeline, gpu_only=False, ignore_less_than_bytes=0):
      """pretty print parsed memory timeline."""
  
      total_memory = 0
      timestamps = []
      data = []
      first_timestamp = timeline[0][0]
      for record in timeline:
          timestamp, kernel_name, allocated_bytes, allocator_type = record
          allocated_bytes = int(allocated_bytes)
  
          if abs(allocated_bytes) < ignore_less_than_bytes:
              continue  # ignore small allocations
  
          total_memory += allocated_bytes
          print(
              "%6d %10d %10d %s"
              % (timestamp - first_timestamp, total_memory, allocated_bytes, kernel_name)
          )
  
  
  def _simplify_device_name(device):
      """/job:localhost/replica:0/task:0/device:CPU:0 -> /cpu:0"""
  
      prefix = '/job:localhost/replica:0/task:0/device:'
      if device.startswith(prefix):
          device = '/' + device[len(prefix) :]
      return device.lower()
  
  
  def _device_stats_dict(run_metadata):
      """Returns dictionary of device_name->[NodeExecStats, NodeExecStats...]"""
      result = {}
      for dev_stat in run_metadata.step_stats.dev_stats:
          device_name = _simplify_device_name(dev_stat.device)
          result[device_name] = dev_stat.node_stats
  
      return result
  
  
  def print_memory_timeline(run_metadata, device=None):
      """Human readable timeline of memory allocation/deallocation for given
    device. If device is None, prints timeline of the device with highest memory
    usage"""
  
      if device is None:
          peak_dict = peak_memory(run_metadata)
          peak_pairs = list(peak_dict.items())
          chosen_peak = _position_of_largest([peak for (dev, peak) in peak_pairs])
          device = peak_pairs[chosen_peak][0]
  
      device_metadata = _device_stats_dict(run_metadata)
      print("Printing timeline for " + device)
      parsed_timeline = _timeline_from_nodestats(device_metadata[device])
      _print_parsed_timeline(parsed_timeline)
  