# 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 the MapAndFilterFusion optimization."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from absl.testing import parameterized

from tensorflow.contrib.data.python.ops import optimization
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import test


class MapAndFilterFusionTest(test.TestCase, parameterized.TestCase):

  @staticmethod
  def map_functions():
    identity = lambda x: x
    increment = lambda x: x + 1

    def increment_and_square(x):
      y = x + 1
      return y * y

    functions = [identity, increment, increment_and_square]
    tests = []
    for i, fun1 in enumerate(functions):
      for j, fun2 in enumerate(functions):
        tests.append((
            "Test{}{}".format(i, j),
            [fun1, fun2],
        ))
        for k, fun3 in enumerate(functions):
          tests.append((
              "Test{}{}{}".format(i, j, k),
              [fun1, fun2, fun3],
          ))

    swap = lambda x, n: (n, x)
    tests.append((
        "Swap1",
        [lambda x: (x, 42), swap],
    ))
    tests.append((
        "Swap2",
        [lambda x: (x, 42), swap, swap],
    ))
    return tuple(tests)

  @parameterized.named_parameters(*map_functions.__func__())
  def testMapFusion(self, functions):
    dataset = dataset_ops.Dataset.range(5).apply(
        optimization.assert_next(["Map", "Prefetch"]))
    for function in functions:
      dataset = dataset.map(function)

    dataset = dataset.prefetch(0).apply(optimization.optimize(["map_fusion"]))
    iterator = dataset.make_one_shot_iterator()
    get_next = iterator.get_next()
    with self.test_session() as sess:
      for x in range(5):
        result = sess.run(get_next)
        r = x
        for function in functions:
          if isinstance(r, tuple):
            r = function(*r)  # Pass tuple as multiple arguments.
          else:
            r = function(r)
        self.assertAllEqual(r, result)

      with self.assertRaises(errors.OutOfRangeError):
        sess.run(get_next)

  @staticmethod
  def map_and_filter_functions():
    identity = lambda x: x
    increment = lambda x: x + 1
    minus_five = lambda x: x - 5

    def increment_and_square(x):
      y = x + 1
      return y * y

    take_all = lambda x: constant_op.constant(True)
    is_zero = lambda x: math_ops.equal(x, 0)
    is_odd = lambda x: math_ops.equal(x % 2, 0)
    greater = lambda x: math_ops.greater(x + 5, 0)

    functions = [identity, increment, minus_five, increment_and_square]
    filters = [take_all, is_zero, is_odd, greater]
    tests = []

    for x, fun in enumerate(functions):
      for y, predicate in enumerate(filters):
        tests.append(("Mixed{}{}".format(x, y), fun, predicate))

    # Multi output
    tests.append(("Multi1", lambda x: (x, x),
                  lambda x, y: constant_op.constant(True)))
    tests.append(
        ("Multi2", lambda x: (x, 2),
         lambda x, y: math_ops.equal(x * math_ops.cast(y, dtypes.int64), 0)))
    return tuple(tests)

  @parameterized.named_parameters(*map_and_filter_functions.__func__())
  def testMapFilterFusion(self, function, predicate):
    dataset = dataset_ops.Dataset.range(10).apply(
        optimization.assert_next(
            ["Map",
             "FilterByLastComponent"])).map(function).filter(predicate).apply(
                 optimization.optimize(["map_and_filter_fusion"]))
    self._testMapAndFilter(dataset, function, predicate)

  def _testMapAndFilter(self, dataset, function, predicate):
    iterator = dataset.make_one_shot_iterator()
    get_next = iterator.get_next()
    with self.test_session() as sess:
      for x in range(10):
        r = function(x)
        if isinstance(r, tuple):
          b = predicate(*r)  # Pass tuple as multiple arguments.
        else:
          b = predicate(r)
        if sess.run(b):
          result = sess.run(get_next)
          self.assertAllEqual(r, result)
      with self.assertRaises(errors.OutOfRangeError):
        sess.run(get_next)

  def testAdditionalInputs(self):
    a = constant_op.constant(3, dtype=dtypes.int64)
    b = constant_op.constant(4, dtype=dtypes.int64)
    some_tensor = math_ops.mul(a, b)
    function = lambda x: x * x

    def predicate(y):
      return math_ops.less(math_ops.cast(y, dtypes.int64), some_tensor)

    # We are currently not supporting functions with additional inputs.
    dataset = dataset_ops.Dataset.range(10).apply(
        optimization.assert_next(
            ["Map", "Filter"])).map(function).filter(predicate).apply(
                optimization.optimize(["map_and_filter_fusion"]))

    self._testMapAndFilter(dataset, function, predicate)

  @staticmethod
  def filter_functions():
    take_all = lambda x: constant_op.constant(True)
    is_zero = lambda x: math_ops.equal(x, 0)
    greater = lambda x: math_ops.greater(x + 5, 0)

    tests = []
    filters = [take_all, is_zero, greater]
    identity = lambda x: x
    for x, predicate_1 in enumerate(filters):
      for y, predicate_2 in enumerate(filters):
        tests.append(("Mixed{}{}".format(x, y), identity,
                      [predicate_1, predicate_2]))
        for z, predicate_3 in enumerate(filters):
          tests.append(("Mixed{}{}{}".format(x, y, z), identity,
                        [predicate_1, predicate_2, predicate_3]))

    take_all_multiple = lambda x, y: constant_op.constant(True)
    # Multi output
    tests.append(("Multi1", lambda x: (x, x),
                  [take_all_multiple, take_all_multiple]))
    tests.append(("Multi2", lambda x: (x, 2), [
        take_all_multiple,
        lambda x, y: math_ops.equal(x * math_ops.cast(y, dtypes.int64), 0)
    ]))
    return tuple(tests)

  @parameterized.named_parameters(*filter_functions.__func__())
  def testFilterFusion(self, map_function, predicates):
    dataset = dataset_ops.Dataset.range(5).apply(
        optimization.assert_next(["Map", "Filter",
                                  "Prefetch"])).map(map_function)
    for predicate in predicates:
      dataset = dataset.filter(predicate)

    dataset = dataset.prefetch(0).apply(
        optimization.optimize(["filter_fusion"]))
    iterator = dataset.make_one_shot_iterator()
    get_next = iterator.get_next()
    with self.test_session() as sess:
      for x in range(5):
        r = map_function(x)
        filtered = False
        for predicate in predicates:
          if isinstance(r, tuple):
            b = predicate(*r)  # Pass tuple as multiple arguments.
          else:
            b = predicate(r)
          if not sess.run(b):
            filtered = True
            break

        if not filtered:
          result = sess.run(get_next)
          self.assertAllEqual(r, result)
      with self.assertRaises(errors.OutOfRangeError):
        sess.run(get_next)


if __name__ == "__main__":
  test.main()

[((224, 2, 224, 13), 'tensorflow.python.platform.test.main', 'test.main', ({}, {}), '()', False, 'from tensorflow.python.platform import test\n'), ((148, 8, 148, 51), 'tensorflow.python.framework.constant_op.constant', 'constant_op.constant', (), '', False, 'from tensorflow.python.framework import constant_op\n'), ((149, 8, 149, 51), 'tensorflow.python.framework.constant_op.constant', 'constant_op.constant', (), '', False, 'from tensorflow.python.framework import constant_op\n'), ((150, 18, 150, 36), 'tensorflow.python.ops.math_ops.mul', 'math_ops.mul', ({(150, 31, 150, 32): 'a', (150, 34, 150, 35): 'b'}, {}), '(a, b)', False, 'from tensorflow.python.ops import math_ops\n'), ((70, 8, 70, 53), 'tensorflow.contrib.data.python.ops.optimization.assert_next', 'optimization.assert_next', ({(70, 33, 70, 52): "['Map', 'Prefetch']"}, {}), "(['Map', 'Prefetch'])", False, 'from tensorflow.contrib.data.python.ops import optimization\n'), ((74, 40, 74, 77), 'tensorflow.contrib.data.python.ops.optimization.optimize', 'optimization.optimize', ({(74, 62, 74, 76): "['map_fusion']"}, {}), "(['map_fusion'])", False, 'from tensorflow.contrib.data.python.ops import optimization\n'), ((101, 25, 101, 51), 'tensorflow.python.framework.constant_op.constant', 'constant_op.constant', ({(101, 46, 101, 50): '(True)'}, {}), '(True)', False, 'from tensorflow.python.framework import constant_op\n'), ((102, 24, 102, 44), 'tensorflow.python.ops.math_ops.equal', 'math_ops.equal', ({(102, 39, 102, 40): 'x', (102, 42, 102, 43): '(0)'}, {}), '(x, 0)', False, 'from tensorflow.python.ops import math_ops\n'), ((103, 23, 103, 47), 'tensorflow.python.ops.math_ops.equal', 'math_ops.equal', ({(103, 38, 103, 43): '(x % 2)', (103, 45, 103, 46): '(0)'}, {}), '(x % 2, 0)', False, 'from tensorflow.python.ops import math_ops\n'), ((104, 24, 104, 50), 'tensorflow.python.ops.math_ops.greater', 'math_ops.greater', ({(104, 41, 104, 46): '(x + 5)', (104, 48, 104, 49): '(0)'}, {}), '(x + 5, 0)', False, 'from tensorflow.python.ops import math_ops\n'), ((128, 17, 128, 65), 'tensorflow.contrib.data.python.ops.optimization.optimize', 'optimization.optimize', ({(128, 39, 128, 64): "['map_and_filter_fusion']"}, {}), "(['map_and_filter_fusion'])", False, 'from tensorflow.contrib.data.python.ops import optimization\n'), ((160, 16, 160, 64), 'tensorflow.contrib.data.python.ops.optimization.optimize', 'optimization.optimize', ({(160, 38, 160, 63): "['map_and_filter_fusion']"}, {}), "(['map_and_filter_fusion'])", False, 'from tensorflow.contrib.data.python.ops import optimization\n'), ((166, 25, 166, 51), 'tensorflow.python.framework.constant_op.constant', 'constant_op.constant', ({(166, 46, 166, 50): '(True)'}, {}), '(True)', False, 'from tensorflow.python.framework import constant_op\n'), ((167, 24, 167, 44), 'tensorflow.python.ops.math_ops.equal', 'math_ops.equal', ({(167, 39, 167, 40): 'x', (167, 42, 167, 43): '(0)'}, {}), '(x, 0)', False, 'from tensorflow.python.ops import math_ops\n'), ((168, 24, 168, 50), 'tensorflow.python.ops.math_ops.greater', 'math_ops.greater', ({(168, 41, 168, 46): '(x + 5)', (168, 48, 168, 49): '(0)'}, {}), '(x + 5, 0)', False, 'from tensorflow.python.ops import math_ops\n'), ((181, 37, 181, 63), 'tensorflow.python.framework.constant_op.constant', 'constant_op.constant', ({(181, 58, 181, 62): '(True)'}, {}), '(True)', False, 'from tensorflow.python.framework import constant_op\n'), ((200, 8, 200, 48), 'tensorflow.contrib.data.python.ops.optimization.optimize', 'optimization.optimize', ({(200, 30, 200, 47): "['filter_fusion']"}, {}), "(['filter_fusion'])", False, 'from tensorflow.contrib.data.python.ops import optimization\n'), ((69, 14, 69, 42), 'tensorflow.python.data.ops.dataset_ops.Dataset.range', 'dataset_ops.Dataset.range', ({(69, 40, 69, 41): '5'}, {}), '(5)', False, 'from tensorflow.python.data.ops import dataset_ops\n'), ((154, 27, 154, 57), 'tensorflow.python.ops.math_ops.cast', 'math_ops.cast', ({(154, 41, 154, 42): 'y', (154, 44, 154, 56): 'dtypes.int64'}, {}), '(y, dtypes.int64)', False, 'from tensorflow.python.ops import math_ops\n'), ((116, 31, 116, 57), 'tensorflow.python.framework.constant_op.constant', 'constant_op.constant', ({(116, 52, 116, 56): '(True)'}, {}), '(True)', False, 'from tensorflow.python.framework import constant_op\n'), ((194, 8, 195, 46), 'tensorflow.contrib.data.python.ops.optimization.assert_next', 'optimization.assert_next', ({(194, 33, 195, 45): "['Map', 'Filter', 'Prefetch']"}, {}), "(['Map', 'Filter', 'Prefetch'])", False, 'from tensorflow.contrib.data.python.ops import optimization\n'), ((193, 14, 193, 42), 'tensorflow.python.data.ops.dataset_ops.Dataset.range', 'dataset_ops.Dataset.range', ({(193, 40, 193, 41): '5'}, {}), '(5)', False, 'from tensorflow.python.data.ops import dataset_ops\n'), ((119, 41, 119, 71), 'tensorflow.python.ops.math_ops.cast', 'math_ops.cast', ({(119, 55, 119, 56): 'y', (119, 58, 119, 70): 'dtypes.int64'}, {}), '(y, dtypes.int64)', False, 'from tensorflow.python.ops import math_ops\n'), ((187, 40, 187, 70), 'tensorflow.python.ops.math_ops.cast', 'math_ops.cast', ({(187, 54, 187, 55): 'y', (187, 57, 187, 69): 'dtypes.int64'}, {}), '(y, dtypes.int64)', False, 'from tensorflow.python.ops import math_ops\n'), ((125, 8, 127, 38), 'tensorflow.contrib.data.python.ops.optimization.assert_next', 'optimization.assert_next', ({(126, 12, 127, 37): "['Map', 'FilterByLastComponent']"}, {}), "(['Map', 'FilterByLastComponent'])", False, 'from tensorflow.contrib.data.python.ops import optimization\n'), ((158, 8, 159, 30), 'tensorflow.contrib.data.python.ops.optimization.assert_next', 'optimization.assert_next', ({(159, 12, 159, 29): "['Map', 'Filter']"}, {}), "(['Map', 'Filter'])", False, 'from tensorflow.contrib.data.python.ops import optimization\n'), ((124, 14, 124, 43), 'tensorflow.python.data.ops.dataset_ops.Dataset.range', 'dataset_ops.Dataset.range', ({(124, 40, 124, 42): '10'}, {}), '(10)', False, 'from tensorflow.python.data.ops import dataset_ops\n'), ((157, 14, 157, 43), 'tensorflow.python.data.ops.dataset_ops.Dataset.range', 'dataset_ops.Dataset.range', ({(157, 40, 157, 42): '10'}, {}), '(10)', False, 'from tensorflow.python.data.ops import dataset_ops\n')]