# 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 vectorization of array kernels."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from tensorflow.python.compat import compat
from tensorflow.python.eager import backprop
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn
from tensorflow.python.ops import random_ops
from tensorflow.python.ops import tensor_array_grad  # pylint: disable=unused-import
from tensorflow.python.ops.parallel_for import control_flow_ops as pfor_control_flow_ops
from tensorflow.python.ops.parallel_for.test_util import PForTestCase
from tensorflow.python.platform import test


@test_util.run_all_in_graph_and_eager_modes
class ArrayTest(PForTestCase):

  def test_gather(self):
    x = random_ops.random_uniform([3, 3, 3])

    def loop_fn(i):
      outputs = []
      x_i = array_ops.gather(x, i)
      for y in [x, x_i]:
        axes = [0, 2, -1] if y is x else [0]
        for axis in axes:
          outputs.append(array_ops.gather(y, 2, axis=axis))
          outputs.append(array_ops.gather(y, i, axis=axis))
          outputs.append(array_ops.gather(y, [i], axis=axis))
          outputs.append(array_ops.gather(y, [i, 2], axis=axis))
          outputs.append(array_ops.gather(y, [[2, i], [i, 1]], axis=axis))
      return outputs

    self._test_loop_fn(loop_fn, 3, loop_fn_dtypes=[dtypes.float32] * 20)

  def test_shape(self):
    x = random_ops.random_uniform([3, 2, 3])

    def loop_fn(i):
      x_i = array_ops.gather(x, i)
      return array_ops.shape(x_i), array_ops.shape(x_i, out_type=dtypes.int64)

    self._test_loop_fn(loop_fn, 3, loop_fn_dtypes=[dtypes.int32, dtypes.int64])

  def test_size(self):
    x = random_ops.random_uniform([3, 2, 3])

    def loop_fn(i):
      x_i = array_ops.gather(x, i)
      return array_ops.size(x_i), array_ops.size(x_i, out_type=dtypes.int64)

    self._test_loop_fn(loop_fn, 3, loop_fn_dtypes=[dtypes.int32, dtypes.int64])

  def test_rank(self):
    x = random_ops.random_uniform([3, 2, 3])

    def loop_fn(i):
      x_i = array_ops.gather(x, i)
      return array_ops.rank(x_i)

    self._test_loop_fn(loop_fn, 3, loop_fn_dtypes=[dtypes.int32])

  def test_shape_n(self):
    x = random_ops.random_uniform([3, 2, 3])
    y = random_ops.random_uniform([3])

    def loop_fn(i):
      x_i = array_ops.gather(x, i)
      y_i = array_ops.gather(y, i)
      return array_ops.shape_n([x_i, x, y, y_i]), array_ops.shape_n(
          [x_i, x, y, y_i], out_type=dtypes.int64)

    self._test_loop_fn(
        loop_fn, 3, loop_fn_dtypes=[dtypes.int32] * 4 + [dtypes.int64] * 4)

  def test_reshape(self):
    x = random_ops.random_uniform([3, 2, 3])

    def loop_fn(i):
      x1 = array_ops.gather(x, i)
      return array_ops.reshape(x1, [-1]), array_ops.reshape(x1, [1, 3, 1, -1])

    self._test_loop_fn(loop_fn, 3, loop_fn_dtypes=[dtypes.float32] * 2)

  def test_broadcast_to(self):
    x = random_ops.random_uniform([3, 2, 1, 3])

    def loop_fn(i):
      x1 = array_ops.gather(x, i)
      return (array_ops.broadcast_to(x1, [2, 2, 3]),
              array_ops.broadcast_to(x1, [1, 2, 1, 3]))

    self._test_loop_fn(loop_fn, 3, loop_fn_dtypes=[dtypes.float32] * 2)

  def test_expand_dims(self):
    x = random_ops.random_uniform([3, 2, 3])

    def loop_fn(i):
      x1 = array_ops.gather(x, i)
      return array_ops.expand_dims(
          x1, axis=-1), array_ops.expand_dims(
              x1, axis=1)

    self._test_loop_fn(loop_fn, 3, loop_fn_dtypes=[dtypes.float32] * 2)

  def test_one_hot(self):
    indices = random_ops.random_uniform(
        [3, 2, 3], minval=0, maxval=4, dtype=dtypes.int32)

    def loop_fn(i):
      indices_i = array_ops.gather(indices, i)
      return (array_ops.one_hot(indices_i, depth=4, on_value=2., off_value=-2.),
              array_ops.one_hot(indices_i, depth=4, axis=1))

    self._test_loop_fn(loop_fn, 3, loop_fn_dtypes=[dtypes.float32] * 2)

  def test_searchsorted(self):
    sorted_inputs = math_ops.cumsum(random_ops.random_uniform([3, 2, 4]),
                                    axis=-1)
    values = random_ops.random_uniform([2, 3], minval=-1, maxval=4.5)

    def loop_fn(i):
      inputs_i = array_ops.gather(sorted_inputs, i)
      return [array_ops.searchsorted(inputs_i, values, out_type=dtypes.int32,
                                     side="left"),  # creates LowerBound op.
              array_ops.searchsorted(inputs_i, values, out_type=dtypes.int64,
                                     side="right")]  # creates UpperBound op.

    self._test_loop_fn(loop_fn, 3, loop_fn_dtypes=[dtypes.int32, dtypes.int64])

  def test_slice(self):
    x = random_ops.random_uniform([3, 2, 3])

    def loop_fn(i):
      x1 = array_ops.gather(x, i)
      return array_ops.slice(x1, begin=(0, 1), size=(2, 1))

    self._test_loop_fn(loop_fn, 3)

  def test_tile(self):
    x = random_ops.random_uniform([3, 2, 3])

    def loop_fn(i):
      x1 = array_ops.gather(x, i)
      return array_ops.tile(x1, [2, 1])

    self._test_loop_fn(loop_fn, 3)

  def test_tile_loop_dependent(self):
    x = random_ops.random_uniform([3, 2, 3])

    def loop_fn(i):
      x1 = array_ops.gather(x, i)
      return array_ops.tile(x1, [i, 1])

    with self.assertRaisesRegexp(ValueError, "expected to be loop invariant"):
      pfor_control_flow_ops.pfor(loop_fn, 2)

  def test_pack(self):
    x = random_ops.random_uniform([3, 2, 3])
    y = random_ops.random_uniform([2, 3])

    def loop_fn(i):
      x1 = array_ops.gather(x, i)
      return array_ops.stack([x1, y], axis=-1)

    self._test_loop_fn(loop_fn, 1)

  def test_unpack(self):
    x = random_ops.random_uniform([3, 2, 3, 4])

    def loop_fn(i):
      x_i = array_ops.gather(x, i)
      return array_ops.unstack(
          x_i, 4, axis=-1), array_ops.unstack(
              x_i, 3, axis=1)

    self._test_loop_fn(loop_fn, 3, loop_fn_dtypes=[dtypes.float32] * 7)

  def test_pad(self):
    x = random_ops.random_uniform([3, 2, 3])
    padding = constant_op.constant([[1, 2], [3, 4]])

    def loop_fn(i):
      x1 = array_ops.gather(x, i)
      return array_ops.pad(x1, padding, mode="CONSTANT")

    self._test_loop_fn(loop_fn, 3)

  def test_split(self):
    x = random_ops.random_uniform([3, 2, 3])

    def loop_fn(i):
      x1 = array_ops.gather(x, i)
      return array_ops.split(x1, 2, axis=0), array_ops.split(x1, 3, axis=-1)

    self._test_loop_fn(loop_fn, 3, loop_fn_dtypes=[dtypes.float32] * 5)

  def test_split_v(self):
    x = random_ops.random_uniform([3, 6, 3])

    def loop_fn(i):
      x1 = array_ops.gather(x, i)
      return (array_ops.split(x1, [2, 1, 3], axis=0),
              array_ops.split(x1, [3], axis=-1))

    self._test_loop_fn(loop_fn, 3, loop_fn_dtypes=[dtypes.float32] * 4)

  def test_squeeze(self):
    x = random_ops.random_uniform([5, 1, 2, 1])

    def loop_fn(i):
      x1 = array_ops.gather(x, i)
      return (array_ops.squeeze(x1, axis=0),
              array_ops.squeeze(x1, axis=-1),
              array_ops.squeeze(x1))

    self._test_loop_fn(loop_fn, 3, loop_fn_dtypes=[dtypes.float32] * 3)

  def test_transpose(self):
    x = random_ops.random_uniform([3, 2, 3, 4])

    def loop_fn(i):
      x1 = array_ops.gather(x, i)
      return array_ops.transpose(x1, [2, 1, 0])

    self._test_loop_fn(loop_fn, 3)

  def test_zeros_like(self):
    x = random_ops.random_uniform([3, 2, 3])

    def loop_fn(i):
      x1 = array_ops.gather(x, i)
      z = array_ops.zeros_like(x1),
      return z, z + x1

    self._test_loop_fn(loop_fn, 3, loop_fn_dtypes=[dtypes.float32] * 2)

  def test_concat_v2(self):
    x = random_ops.random_uniform([3, 2, 3])
    y = random_ops.random_uniform([2, 3])

    def loop_fn(i):
      x1 = array_ops.gather(x, i)
      return array_ops.concat(
          [x1, x1, y], axis=0), array_ops.concat(
              [x1, x1, y], axis=-1)

    self._test_loop_fn(loop_fn, 3, loop_fn_dtypes=[dtypes.float32] * 2)

  def test_unary_cwise_ops(self):
    for op in [array_ops.identity, array_ops.stop_gradient]:
      with backprop.GradientTape(persistent=True) as g:
        x = random_ops.random_uniform([3, 5])
        g.watch(x)

      # pylint: disable=cell-var-from-loop
      def loop_fn(i):
        with g:
          x1 = array_ops.gather(x, i)
          y = op(x1) + x1
          loss = nn.l2_loss(y)
        return op(x), y, g.gradient(loss, x1)

      # pylint: enable=cell-var-from-loop

      self._test_loop_fn(loop_fn, 3, loop_fn_dtypes=[dtypes.float32] * 3)

  def test_identity_n(self):
    x = random_ops.random_uniform([3, 4])

    def loop_fn(i):
      return array_ops.identity_n([x, array_ops.gather(x, i)])

    self._test_loop_fn(loop_fn, 3, loop_fn_dtypes=[dtypes.float32] * 2)

  def test_matrix_band_part(self):
    x = random_ops.random_uniform([3, 4, 2, 2])

    for num_lower, num_upper in ((0, -1), (-1, 0), (1, 1)):
      # pylint: disable=cell-var-from-loop
      def loop_fn(i):
        return array_ops.matrix_band_part(
            array_ops.gather(x, i),
            num_lower=num_lower,
            num_upper=num_upper)
      # pylint: enable=cell-var-from-loop

    self._test_loop_fn(loop_fn, 3)

  def test_matrix_diag(self):
    x = random_ops.random_uniform([3, 2, 4])

    def loop_fn(i):
      diagonal = array_ops.gather(x, i)
      if compat.forward_compatible(2019, 8, 31):
        return array_ops.matrix_diag(diagonal, k=(0, 1), num_rows=4, num_cols=5)
      return array_ops.matrix_diag(diagonal)

    self._test_loop_fn(loop_fn, 3, loop_fn_dtypes=[dtypes.float32])

  def test_matrix_diag_part(self):
    x = random_ops.random_uniform([3, 4, 6])

    def loop_fn(i):
      input = array_ops.gather(x, i)  # pylint: disable=redefined-builtin
      if compat.forward_compatible(2019, 8, 31):
        return array_ops.matrix_diag_part(input, k=(-2, 0), padding_value=3)
      return array_ops.matrix_diag_part(input)

    self._test_loop_fn(loop_fn, 3, loop_fn_dtypes=[dtypes.float32])

  def test_matrix_set_diag(self):
    matrices = random_ops.random_uniform([3, 4, 4])
    diags = random_ops.random_uniform([3, 4])
    num_outputs = 3
    if compat.forward_compatible(2019, 8, 31):
      bands = random_ops.random_uniform([3, 3, 4])
      num_outputs = 6

    def loop_fn(i):
      matrix_i = array_ops.gather(matrices, i)
      diag_i = array_ops.gather(diags, i)
      results = [
          array_ops.matrix_set_diag(matrix_i, diag_i),
          array_ops.matrix_set_diag(matrices[0, ...], diag_i),
          array_ops.matrix_set_diag(matrix_i, diags[0, ...])
      ]
      if compat.forward_compatible(2019, 8, 31):
        k = (-1, 1)
        band_i = array_ops.gather(bands, i)
        results.extend([
            array_ops.matrix_set_diag(matrix_i, band_i, k=k),
            array_ops.matrix_set_diag(matrices[0, ...], band_i, k=k),
            array_ops.matrix_set_diag(matrix_i, bands[0, ...], k=k)
        ])
      return results

    self._test_loop_fn(
        loop_fn, 3, loop_fn_dtypes=[dtypes.float32] * num_outputs)

  def test_strided_slice(self):
    with backprop.GradientTape(persistent=True) as g:
      x = random_ops.random_uniform([3, 3, 4, 4, 2, 2, 2])
      g.watch(x)

    def loop_fn(i):
      with g:
        x_i = array_ops.gather(x, i)
        y = x_i[:2, ::2, 1::3, ..., array_ops.newaxis, 1]
        loss = nn.l2_loss(y)
      return y, g.gradient(loss, x_i)

    self._test_loop_fn(loop_fn, 3, loop_fn_dtypes=[dtypes.float32] * 2)


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