xla_device_test.py 文件源码

def testLoops(self):
    """Tests that loops work on XLA devices."""

    with session_lib.Session() as session:
      x = array_ops.placeholder(dtypes.float32)
      with ops.device("device:XLA_CPU:0"):
        c = lambda i, _: math_ops.less(i, 5)
        b = lambda i, x: (i + 1, x * 2.0 + 1.0)
        _, y = control_flow_ops.while_loop(c, b, (constant_op.constant(0), x))

      result = session.run(y, {x: np.float32(2)})
      self.assertAllClose(result, np.float32(95), rtol=1e-3)