variable_ops_test.py 文件源码

def testAssignDependencyAcrossDevices(self):
    with self.test_session(use_gpu=True):
      # The variable and an op to increment it are on the GPU.
      var = state_ops.variable_op([1], tf.complex64)
      tf.assign(var, [1.0+2.0j]).eval()
      increment = tf.assign_add(var, [2.0+3.0j])
      with tf.control_dependencies([increment]):
        with tf.device("/cpu:0"):
          # This mul op is pinned to the CPU, but reads the variable from the
          # GPU. The test ensures that the dependency on 'increment' is still
          # honored, i.e., the Send and Recv from GPU to CPU should take place
          # only after the increment.
          result = tf.multiply(var, var)
      self.assertAllClose([-16.0+30.0j], result.eval())