tensorflow/compiler/tests/while_test.py

# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for while loops in XLA."""

import os

import numpy as np

from tensorflow.compiler.tests import xla_test
from tensorflow.compiler.tf2xla.python import xla
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import function
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import gradients_impl
from tensorflow.python.ops import map_fn
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import while_loop
from tensorflow.python.platform import test


class WhileTest(xla_test.XLATestCase):

  def testSingletonLoopHandrolled(self):
    # Define a function for the loop body
    @function.Defun(dtypes.int32)
    def loop_body(step):
      step_out = step + constant_op.constant(1, dtype=dtypes.int32)
      return step_out

    # Define a function for the loop condition
    @function.Defun(dtypes.int32)
    def loop_cond(step):
      return step < 10

    with self.session() as sess:
      init_index = array_ops.placeholder(dtypes.int32, [])
      with self.test_scope():
        loop_outputs = xla.while_loop([init_index], loop_cond, loop_body)

      result = sess.run(loop_outputs, {init_index: 0})
      self.assertAllClose(result, [10], rtol=1e-3)

  def testCountingLoopHandrolled(self):
    # Define a function for the loop body
    @function.Defun(dtypes.int32, dtypes.float32)
    def loop_body(step, rsum):
      step_out = step + constant_op.constant(1, dtype=dtypes.int32)
      sum_out = rsum + constant_op.constant(1.5, dtype=dtypes.float32)
      return step_out, sum_out

    # Define a function for the loop condition
    @function.Defun(dtypes.int32, dtypes.float32)
    def loop_cond(step, rsum):
      del rsum
      return step < 10

    with self.session() as sess:
      init_index = array_ops.placeholder(dtypes.int32, [])
      init_sum = array_ops.placeholder(dtypes.float32, [])
      with self.test_scope():
        loop_outputs = xla.while_loop([init_index, init_sum], loop_cond,
                                      loop_body)

      result = sess.run(loop_outputs, {init_index: 0, init_sum: 0.0})
      self.assertAllClose(result, [10, 15.0], rtol=1e-3)
      no_iters_result = sess.run(loop_outputs, {init_index: 10, init_sum: 0.0})
      self.assertAllClose(no_iters_result, [10, 0.0], rtol=1e-3)

  def testCountingLoopHandrolledC64(self):
    # Define a function for the loop body
    @function.Defun(dtypes.int32, dtypes.complex64)
    def loop_body(step, rsum):
      step_out = step + constant_op.constant(1, dtype=dtypes.int32)
      sum_out = rsum + constant_op.constant(1.5 + 2j, dtype=dtypes.complex64)
      return step_out, sum_out

    # Define a function for the loop condition
    @function.Defun(dtypes.int32, dtypes.complex64)
    def loop_cond(step, rsum):
      del rsum
      return step < 10

    with self.session() as sess:
      init_index = array_ops.placeholder(dtypes.int32, [])
      init_sum = array_ops.placeholder(dtypes.complex64, [])
      with self.test_scope():
        loop_outputs = xla.while_loop([init_index, init_sum], loop_cond,
                                      loop_body)

      result = sess.run(loop_outputs, {init_index: 0, init_sum: 0.0})
      self.assertAllClose(result[1], np.complex64(15 + 20j), rtol=1e-3)
      no_iters_result = sess.run(loop_outputs, {init_index: 10, init_sum: 0.0})
      self.assertAllClose(no_iters_result[1], np.complex64(0), rtol=1e-3)

  def testLoopWithConstantOutput(self):
    # Define a function for the loop body
    @function.Defun(dtypes.int32, dtypes.int32)
    def loop_body(step, x):
      del x
      step_out = step + constant_op.constant(1, dtype=dtypes.int32)
      return (step_out, 7)

    # Define a function for the loop condition
    @function.Defun(dtypes.int32, dtypes.int32)
    def loop_cond(step, x):
      del x
      return step < 10

    with self.session() as sess:
      init_index = array_ops.placeholder(dtypes.int32, [])
      with self.test_scope():
        loop_outputs = xla.while_loop([init_index, 42], loop_cond, loop_body)

      result = sess.run(loop_outputs, {init_index: 0})
      self.assertAllClose(result, [10, 7], rtol=1e-3)

  def _testMaxItersSimple(self):
    if is_compile_on_demand():
      self.skipTest("list_ops are not supported in cpu_ondemand")
    with self.session() as sess, self.test_scope():
      xla_context = control_flow_ops.XLAControlFlowContext()
      xla_context.Enter()
      v = constant_op.constant(1.0)
      p = array_ops.placeholder(dtype=dtypes.int32)

      def create_while_loop():
        iterations = array_ops.size(p, name="iterations")
        r = while_loop.while_loop(
            lambda *_: True,
            lambda i, x: (i + 1, v * x), (0, 1.0),
            maximum_iterations=iterations,
            name="outer")
        return array_ops.identity(r[1])

      output = create_while_loop()
      output = gradients_impl.gradients(output, v)[0]

      result = sess.run(output, feed_dict={p: [0, 0, 0]})
      print(result)
      xla_context.Exit()

  def testMaxItersSimple(self):
    self.skipTest("Fails with v1 control flow")
    # This fails with old control.
    # self._testMaxItersSimple()

  @test_util.enable_control_flow_v2
  def testMaxItersSimpleV2(self):
    self._testMaxItersSimple()

  def _testNestedWhileLoopWithMaxItersFromOuterContext(self):
    if is_compile_on_demand():
      self.skipTest("list_ops are not supported in cpu_ondemand")
    with self.session() as sess, self.test_scope():
      xla_context = control_flow_ops.XLAControlFlowContext()
      xla_context.Enter()
      v = constant_op.constant(1.0)
      p = array_ops.placeholder(dtype=dtypes.int32)

      def mid_body_builder(iterations):

        def mid_body(i, x):
          r = while_loop.while_loop(
              lambda *_: True,
              lambda i, x: (i + 1, v * x), (0, x),
              maximum_iterations=iterations,
              name="inner")
          return (i + 1, gradients_impl.gradients(x + r[1], v)[0])

        return mid_body

      def outer_body(i, x):
        iterations = array_ops.size(p, name="iterations")
        return (i + 1, x + while_loop.while_loop(
            lambda *_: True,
            mid_body_builder(iterations), (0, x),
            maximum_iterations=iterations,
            name="mid")[1])

      def create_while_loop():
        r = while_loop.while_loop(
            lambda *_: True,
            outer_body, (0, 1.0),
            maximum_iterations=5,
            name="outer")
        return array_ops.identity(r[1])

      # p:placeholder
      # j = 0
      # i, x = 0, 1.
      # while j++ < 5:
      #   i1, x1 = 0, x
      #   while i1++ < len(p):
      #     i2, x2 = 0, x1
      #     while i2++ < len(p):
      #       x2 = v * x2
      #     x1 = grad(x1 + x2, v)
      #   x = x1
      # output = x
      output = create_while_loop()
      sess.run(output, feed_dict={p: [0, 0, 0]})
      xla_context.Exit()

  def testNestedWhileLoopWithMaxItersFromOuterContext(self):
    self._testNestedWhileLoopWithMaxItersFromOuterContext()

  @test_util.enable_control_flow_v2
  def testNestedWhileLoopWithMaxItersFromOuterContextV2(self):
    self._testNestedWhileLoopWithMaxItersFromOuterContext()

  @test_util.enable_control_flow_v2
  def testMap(self):
    if is_compile_on_demand():
      self.skipTest("list_ops are not supported in cpu_ondemand")
    with self.session(), self.test_scope():
      xla_context = control_flow_ops.XLAControlFlowContext()
      xla_context.Enter()
      nums = [1, 2, 3, 4, 5, 6]
      elems = constant_op.constant(nums, name="data")
      r = map_fn.map_fn(lambda x: math_ops.multiply(math_ops.add(x, 3), 2),
                        elems)
      self.assertAllEqual(r, np.array([(x + 3) * 2 for x in nums]))
      xla_context.Exit()

  @test_util.enable_control_flow_v2
  def testMapBackPropFalse(self):
    if is_compile_on_demand():
      self.skipTest("list_ops are not supported in cpu_ondemand")
    with self.session(), self.test_scope():
      xla_context = control_flow_ops.XLAControlFlowContext()
      xla_context.Enter()
      nums = [1, 2, 3, 4, 5, 6]
      elems = constant_op.constant(nums, name="data")
      r = map_fn.map_fn(
          lambda x: math_ops.multiply(math_ops.add(x, 3), 2),
          elems,
          back_prop=False)
      self.assertAllEqual(r, np.array([(x + 3) * 2 for x in nums]))
      xla_context.Exit()


def is_compile_on_demand():
  return ("TF_XLA_FLAGS" in os.environ and
          "tf_xla_compile_on_demand" in os.environ["TF_XLA_FLAGS"])


if __name__ == "__main__":
  os.environ["TF_XLA_FLAGS"] = ("--tf_xla_min_cluster_size=2 " +
                                os.environ.get("TF_XLA_FLAGS", ""))
  test.main()