class_docstr: Place the variable on different device and balance the memory usage.
class_docstr_tok: ["Place","the","variable","on","different","device","and","balance","the","memory","usage","."]
code: def __init__(self, devices, init_usage=None):
      init_usage = init_usage if init_usage else [0] * len(devices)
      assert len(devices) == len(init_usage)
      self._mem_device_heap = list(zip(init_usage, devices))
      heapq.heapify(self._mem_device_heap)
      self._last_device = devices[0]
  
file_path: mesh/mesh_tensorflow/utils.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import contextlib
  import heapq
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow.python.framework import ops
  from tensorflow.python.tpu.topology import Topology
label_desc: Use image node_id 1 to create a new BalancedVariablePlacer object from inherited base classes: object with example: obj = BalancedVariablePlacer(devices, init_usage)