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","."]
docstr: Choose a device for the input variable.
  
  Args:
    var: an Variable.
  
  Returns:
    The device for placing the var.
docstr_tok: ["Choose","a","device","for","the","input","variable",".","Args",":","var",":","an","Variable",".","Returns",":","The","device","for","placing","the","var","."]
code: def device_function(self, var):
      """Choose a device for the input variable.
  
      Args:
        var: an Variable.
  
      Returns:
        The device for placing the var.
      """
      if var.type not in ("Variable", "VariableV2", "VarHandleOp"):
          tf.logging.debug(
              "Place {} on last device: {}.".format(
                  var.name, self._last_device
              )
          )
          return self._last_device
  
      shape = tf.TensorShape(var.get_attr("shape"))
      assert shape.num_elements() is not None
  
      size = var.get_attr("dtype").size
      mem, device = heapq.heappop(self._mem_device_heap)
      mem += shape.num_elements() * size
      heapq.heappush(self._mem_device_heap, (mem, device))
      tf.logging.debug(
          "Place variable {} on {} and consumes {} Bytes.".format(
              var.name, device, mem
          )
      )
      self._last_device = device
  
      return device
  
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 2 for calling the BalancedVariablePlacer obj's underlying member method code with example usage: obj.device_function(var) and returns: device, self