# Copyright 2015 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 tensorflow.ops.tf.gather."""

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

import numpy as np

from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gradients_impl
from tensorflow.python.platform import test

_TEST_TYPES = (dtypes.float32, dtypes.complex64, dtypes.complex128)


class GatherTest(test.TestCase):

  def _buildParams(self, data, dtype):
    data = data.astype(dtype.as_numpy_dtype)
    # For complex types, add an index-dependent imaginary component so we can
    # tell we got the right value.
    if dtype.is_complex:
      return data + 10j * data
    return data

  def testScalar1D(self):
    with self.test_session(use_gpu=True):
      data = np.array([0, 1, 2, 3, 7, 5])
      for dtype in _TEST_TYPES:
        for indices in 4, [1, 2, 2, 4, 5]:
          params_np = self._buildParams(data, dtype)
          params = constant_op.constant(params_np)
          indices_tf = constant_op.constant(indices)
          gather_t = array_ops.gather(params, indices_tf)
          gather_val = gather_t.eval()
          np_val = params_np[indices]
          self.assertAllEqual(np_val, gather_val)
          self.assertEqual(np_val.shape, gather_t.get_shape())

  def testScalar2D(self):
    with self.test_session(use_gpu=True):
      data = np.array([[0, 1, 2], [3, 4, 5], [6, 7, 8],
                       [9, 10, 11], [12, 13, 14]])
      for dtype in _TEST_TYPES:
        for axis in range(data.ndim):
          params_np = self._buildParams(data, dtype)
          params = constant_op.constant(params_np)
          indices = constant_op.constant(2)
          gather_t = array_ops.gather(params, indices, axis=axis)
          gather_val = gather_t.eval()
          self.assertAllEqual(np.take(params_np, 2, axis=axis), gather_val)
          expected_shape = data.shape[:axis] + data.shape[axis + 1:]
          self.assertEqual(expected_shape, gather_t.get_shape())

  def testSimpleTwoD32(self):
    with self.test_session(use_gpu=True):
      data = np.array([[0, 1, 2], [3, 4, 5], [6, 7, 8],
                       [9, 10, 11], [12, 13, 14]])
      for dtype in _TEST_TYPES:
        for axis in range(data.ndim):
          params_np = self._buildParams(data, dtype)
          params = constant_op.constant(params_np)
          # The indices must be in bounds for any axis.
          indices = constant_op.constant([0, 1, 0, 2])
          gather_t = array_ops.gather(params, indices, axis=axis)
          gather_val = gather_t.eval()
          self.assertAllEqual(np.take(params_np, [0, 1, 0, 2], axis=axis),
                              gather_val)
          expected_shape = data.shape[:axis] + (4,) + data.shape[axis + 1:]
          self.assertEqual(expected_shape, gather_t.get_shape())

  def testHigherRank(self):
    # We check that scalar and empty indices shapes work as well
    for shape in (4, 3, 2), (2, 1, 3, 2):
      for indices_shape in (), (0,), (3, 0), (3, 5), (5, 2, 3):
        for dtype in _TEST_TYPES:
          for axis in range(len(shape)):
            params = self._buildParams(np.random.randn(*shape), dtype)
            indices = np.random.randint(shape[axis], size=indices_shape)
            with self.test_session(use_gpu=True) as sess:
              tf_params = constant_op.constant(params)
              tf_indices = constant_op.constant(indices)
              # Check that both positive and negative indices for axis work.
              tf_axis = constant_op.constant(axis)
              tf_negative_axis = constant_op.constant(-len(shape) + axis)
              gather = array_ops.gather(tf_params, tf_indices, axis=tf_axis)
              gather_negative_axis = array_ops.gather(
                  tf_params, tf_indices, axis=tf_negative_axis)
              gather_value, gather_negative_axis_value = sess.run(
                  [gather, gather_negative_axis])
              gather_np = np.take(params, indices, axis)
              self.assertAllEqual(gather_np, gather_value)
              self.assertAllEqual(gather_np, gather_negative_axis_value)
              expected_shape = (params.shape[:axis] + indices.shape +
                                params.shape[axis + 1:])
              self.assertEqual(expected_shape, gather.shape)
              self.assertEqual(expected_shape, gather_negative_axis.shape)

              # Test gradients
              gather_grad = np.random.randn(
                  *gather.get_shape().as_list()).astype(dtype.as_numpy_dtype)
              if dtype.is_complex:
                gather_grad -= 1j * gather_grad
              params_grad, indices_grad, axis_grad = gradients_impl.gradients(
                  gather, [tf_params, tf_indices, tf_axis], gather_grad)
              self.assertEqual(indices_grad, None)
              self.assertEqual(axis_grad, None)
              # For axis 0, we are able to create an efficient IndexedSlices for
              # the gradient.
              if axis == 0:
                self.assertEqual(type(params_grad), ops.IndexedSlices)
                params_grad = ops.convert_to_tensor(params_grad)
              correct_params_grad = np.zeros(shape).astype(dtype.as_numpy_dtype)
              outer_dims = axis
              inner_dims = len(shape) - axis - 1
              gather_grad = gather_grad.reshape(
                  shape[:axis] + (indices.size,) + shape[axis + 1:])
              for source_index, dest_index in enumerate(indices.flat):
                dest_slice = ((slice(None),) * outer_dims + (dest_index,) +
                              (slice(None),) * inner_dims)
                source_slice = ((slice(None),) * outer_dims + (source_index,) +
                                (slice(None),) * inner_dims)
                correct_params_grad[dest_slice] += gather_grad[source_slice]
              self.assertAllClose(correct_params_grad, params_grad.eval(),
                                  atol=2e-6, rtol=2e-6)

  def testString(self):
    params = np.array([[b"asdf", b"zxcv"], [b"qwer", b"uiop"]])
    with self.test_session():
      self.assertAllEqual([b"qwer", b"uiop"],
                          array_ops.gather(params, 1, axis=0).eval())
      self.assertAllEqual([b"asdf", b"qwer"],
                          array_ops.gather(params, 0, axis=1).eval())

  def testUnknownIndices(self):
    params = constant_op.constant([[0, 1, 2]])
    indices = array_ops.placeholder(dtypes.int32)
    gather_t = array_ops.gather(params, indices)
    self.assertEqual(None, gather_t.get_shape())

  def testUnknownAxis(self):
    params = constant_op.constant([[0, 1, 2]])
    indices = constant_op.constant([[0, 0], [0, 0]])
    axis = array_ops.placeholder(dtypes.int32)
    gather_t = array_ops.gather(params, indices, axis=axis)
    # Rank 2 params with rank 2 indices results in a rank 3 shape.
    self.assertEqual([None, None, None], gather_t.shape.as_list())

    # If indices is also unknown the result rank is unknown.
    indices = array_ops.placeholder(dtypes.int32)
    gather_t = array_ops.gather(params, indices, axis=axis)
    self.assertEqual(None, gather_t.shape)

  def testBadIndices(self):
    with self.test_session(use_gpu=True):
      params = [[0, 1, 2], [3, 4, 5]]
      with self.assertRaisesOpError(r"indices\[0,0\] = 7 is not in \[0, 2\)"):
        array_ops.gather(params, [[7]], axis=0).eval()
      with self.assertRaisesOpError(r"indices\[0,0\] = 7 is not in \[0, 3\)"):
        array_ops.gather(params, [[7]], axis=1).eval()

  def testBadAxis(self):
    with self.test_session(use_gpu=True):
      params = [0, 1, 2]
      params_ph = array_ops.placeholder(dtypes.int32)
      indices = 0
      for bad_axis in (1, 2, -2):
        # Shape inference can validate axis for known params rank.
        with self.assertRaisesWithPredicateMatch(
            ValueError, "Shape must be at least rank . but is rank 1"):
          array_ops.gather(params, indices, axis=bad_axis)
        # If params rank is unknown, an op error occurs.
        with self.assertRaisesOpError(
            r"Expected axis in the range \[-1, 1\), but got %s" % bad_axis):
          array_ops.gather(params_ph, indices, axis=bad_axis).eval(
              feed_dict={params_ph: params})

  def testEmptySlices(self):
    with self.test_session(use_gpu=True):
      for dtype in _TEST_TYPES:
        for itype in np.int32, np.int64:
          # Leading axis gather.
          params = np.zeros((7, 0, 0), dtype=dtype.as_numpy_dtype)
          indices = np.array([3, 4], dtype=itype)
          gather = array_ops.gather(params, indices, axis=0)
          self.assertAllEqual(gather.eval(), np.zeros((2, 0, 0)))

          # Middle axis gather.
          params = np.zeros((0, 7, 0), dtype=dtype.as_numpy_dtype)
          gather = array_ops.gather(params, indices, axis=1)
          self.assertAllEqual(gather.eval(), np.zeros((0, 2, 0)))

          # Trailing axis gather.
          params = np.zeros((0, 0, 7), dtype=dtype.as_numpy_dtype)
          gather = array_ops.gather(params, indices, axis=2)
          self.assertAllEqual(gather.eval(), np.zeros((0, 0, 2)))


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