<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 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