# 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 SparseTensorsMap."""

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

import numpy as np

from tensorflow.python.client import session
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import sparse_tensor as sparse_tensor_lib
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import sparse_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import benchmark
from tensorflow.python.platform import test

# pylint: disable=protected-access
add_sparse_to_tensors_map = sparse_ops._add_sparse_to_tensors_map
add_many_sparse_to_tensors_map = sparse_ops._add_many_sparse_to_tensors_map
take_many_sparse_from_tensors_map = (
    sparse_ops._take_many_sparse_from_tensors_map)

# pylint: enable=protected-access


class SparseTensorsMapTest(test.TestCase):

  def _SparseTensorPlaceholder(self, dtype=None):
    if dtype is None:
      dtype = dtypes.int32
    return sparse_tensor_lib.SparseTensor(
        array_ops.placeholder(dtypes.int64),
        array_ops.placeholder(dtype), array_ops.placeholder(dtypes.int64))

  def _SparseTensorValue_5x6(self, permutation):
    ind = np.array([[0, 0], [1, 0], [1, 3], [1, 4], [3, 2],
                    [3, 3]]).astype(np.int64)
    val = np.array([0, 10, 13, 14, 32, 33]).astype(np.int32)

    ind = ind[permutation]
    val = val[permutation]

    shape = np.array([5, 6]).astype(np.int64)
    return sparse_tensor_lib.SparseTensorValue(ind, val, shape)

  def _SparseTensorValue_3x4(self, permutation):
    ind = np.array([[0, 0], [1, 0], [1, 2], [1, 3], [2, 2],
                    [2, 3]]).astype(np.int64)
    val = np.array([0, 10, 13, 14, 32, 33]).astype(np.int32)

    ind = ind[permutation]
    val = val[permutation]

    shape = np.array([3, 4]).astype(np.int64)
    return sparse_tensor_lib.SparseTensorValue(ind, val, shape)

  def _SparseTensorValue_1x1x1(self):
    ind = np.array([[0, 0, 0]]).astype(np.int64)
    val = np.array([0]).astype(np.int32)
    shape = np.array([3, 4, 5]).astype(np.int64)
    return sparse_tensor_lib.SparseTensorValue(ind, val, shape)

  def testAddTakeMany(self):
    with self.session(graph=ops.Graph(), use_gpu=False) as sess:
      sp_input0 = self._SparseTensorValue_5x6(np.arange(6))
      sp_input1 = self._SparseTensorValue_3x4(np.arange(6))
      handle0 = add_sparse_to_tensors_map(sp_input0, shared_name="a")
      handle1 = add_sparse_to_tensors_map(sp_input1, shared_name="a")
      self.assertEqual(handle0.get_shape(), ())
      handles_concat = array_ops.stack([handle0, handle1])

      sp_out = take_many_sparse_from_tensors_map(
          sparse_map_op=handle0.op, sparse_handles=handles_concat)

      combined_indices, combined_values, combined_shape = self.evaluate(sp_out)

      self.assertAllEqual(combined_indices[:6, 0], [0] * 6)  # minibatch 0
      self.assertAllEqual(combined_indices[:6, 1:], sp_input0[0])
      self.assertAllEqual(combined_indices[6:, 0], [1] * 6)  # minibatch 1
      self.assertAllEqual(combined_indices[6:, 1:], sp_input1[0])
      self.assertAllEqual(combined_values[:6], sp_input0[1])
      self.assertAllEqual(combined_values[6:], sp_input1[1])
      self.assertAllEqual(combined_shape, [2, 5, 6])

  def testFeedAddTakeMany(self):
    with self.session(use_gpu=False) as sess:
      sp_input = self._SparseTensorPlaceholder()
      input0_val = self._SparseTensorValue_5x6(np.arange(6))
      input1_val = self._SparseTensorValue_3x4(np.arange(6))
      handle = add_sparse_to_tensors_map(sp_input)

      handle0_value = sess.run(handle, feed_dict={sp_input: input0_val})
      handle1_value = sess.run(handle, feed_dict={sp_input: input1_val})

      sparse_handles = ops.convert_to_tensor(
          [handle0_value, handle1_value], dtype=dtypes.int64)

      sp_roundtrip = take_many_sparse_from_tensors_map(
          sparse_map_op=handle.op, sparse_handles=sparse_handles)

      combined_indices, combined_values, combined_shape = self.evaluate(
          sp_roundtrip)

      self.assertAllEqual(combined_indices[:6, 0], [0] * 6)  # minibatch 0
      self.assertAllEqual(combined_indices[:6, 1:], input0_val[0])
      self.assertAllEqual(combined_indices[6:, 0], [1] * 6)  # minibatch 1
      self.assertAllEqual(combined_indices[6:, 1:], input1_val[0])
      self.assertAllEqual(combined_values[:6], input0_val[1])
      self.assertAllEqual(combined_values[6:], input1_val[1])
      self.assertAllEqual(combined_shape, [2, 5, 6])

  def testAddManyTakeManyRoundTrip(self):
    with self.session(use_gpu=False) as sess:
      # N == 4 because shape_value == [4, 5]
      indices_value = np.array([[0, 0], [0, 1], [2, 0]], dtype=np.int64)
      values_value = np.array([b"a", b"b", b"c"])
      shape_value = np.array([4, 5], dtype=np.int64)
      sparse_tensor = self._SparseTensorPlaceholder(dtype=dtypes.string)
      handles = add_many_sparse_to_tensors_map(sparse_tensor)
      roundtrip = take_many_sparse_from_tensors_map(
          sparse_map_op=handles.op, sparse_handles=handles)
      handles_value, roundtrip_value = sess.run(
          [handles, roundtrip],
          feed_dict={
              sparse_tensor.indices: indices_value,
              sparse_tensor.values: values_value,
              sparse_tensor.dense_shape: shape_value
          })
      self.assertEqual(handles_value.shape, (4,))
      self.assertAllEqual(roundtrip_value.indices, indices_value)
      self.assertAllEqual(roundtrip_value.values, values_value)
      self.assertAllEqual(roundtrip_value.dense_shape, shape_value)

  def testDeserializeFailsInconsistentRank(self):
    with self.session(use_gpu=False) as sess:
      sp_input = self._SparseTensorPlaceholder()
      input0_val = self._SparseTensorValue_5x6(np.arange(6))
      input1_val = self._SparseTensorValue_1x1x1()
      handle = add_sparse_to_tensors_map(sp_input)

      handle0_value = sess.run(handle, feed_dict={sp_input: input0_val})
      handle1_value = sess.run(handle, feed_dict={sp_input: input1_val})

      handle_concat = ops.convert_to_tensor(
          [handle0_value, handle1_value], dtype=dtypes.int64)

      sp_roundtrip = take_many_sparse_from_tensors_map(
          sparse_map_op=handle.op, sparse_handles=handle_concat)

      with self.assertRaisesOpError(
          r"Inconsistent rank across SparseTensors: rank prior to "
          r"SparseTensor\[1\] was: 3 but rank of SparseTensor\[1\] is: 4"):
        self.evaluate(sp_roundtrip)

  def testTakeManyFailsWrongInputOp(self):
    with self.session(use_gpu=False) as sess:
      input_val = self._SparseTensorValue_5x6(np.arange(6))
      handle = add_sparse_to_tensors_map(input_val)
      handle_value = self.evaluate(handle)
      bad_handle = handle_value + 10
      sp_roundtrip = take_many_sparse_from_tensors_map(
          sparse_map_op=handle.op, sparse_handles=[handle_value, bad_handle])

      with self.assertRaisesOpError(r"Unable to find SparseTensor: 10"):
        self.evaluate(sp_roundtrip)


class BenchmarkSparseTensorsMapVsSerialization(test.Benchmark):

  def benchmarkVeryLarge2DFloatSparseTensor(self):
    np.random.seed(127)
    num_elements = 10000
    batch_size = 64
    indices_batch = np.random.randint(
        batch_size, size=num_elements, dtype=np.int64)
    indices_value = np.arange(num_elements, dtype=np.int64)
    indices = np.asarray(
        sorted(zip(indices_batch, indices_value)), dtype=np.int64)
    values = ["feature_value_for_embedding_lookup"] * num_elements
    shape = np.asarray([batch_size, num_elements], dtype=np.int64)
    with session.Session(config=benchmark.benchmark_config()) as sess:
      with ops.device("/cpu:0"):
        indices = variables.Variable(indices)
        values = variables.Variable(values)
        shape = variables.Variable(shape)
        st = sparse_tensor_lib.SparseTensor(indices, values, shape)

        st_handles = add_many_sparse_to_tensors_map(st)
        st_roundtrip = take_many_sparse_from_tensors_map(
            sparse_map_op=st_handles.op, sparse_handles=st_handles)
        st_roundtrip_op = st_roundtrip.values.op

        st_serialized = sparse_ops.serialize_many_sparse(st)
        st_deserialized = sparse_ops.deserialize_many_sparse(
            st_serialized, dtype=values.dtype)
        st_deserialized_op = st_deserialized.values.op

        variables.global_variables_initializer().run()

        st_roundtrip_values = self.evaluate(st_roundtrip)
        st_deserialized_values = self.evaluate(st_deserialized)
        np.testing.assert_equal(st_roundtrip_values.values,
                                st_deserialized_values.values)
        np.testing.assert_equal(st_roundtrip_values.indices,
                                st_deserialized_values.indices)
        np.testing.assert_equal(st_roundtrip_values.dense_shape,
                                st_deserialized_values.dense_shape)

        self.run_op_benchmark(
            sess,
            st_roundtrip_op,
            min_iters=2000,
            name="benchmark_very_large_2d_float_st_tensor_maps")
        self.run_op_benchmark(
            sess,
            st_deserialized_op,
            min_iters=2000,
            name="benchmark_very_large_2d_float_st_serialization")


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

[((238, 2, 238, 13), 'tensorflow.python.platform.test.main', 'test.main', ({}, {}), '()', False, 'from tensorflow.python.platform import test\n'), ((60, 11, 60, 63), 'tensorflow.python.framework.sparse_tensor.SparseTensorValue', 'sparse_tensor_lib.SparseTensorValue', ({(60, 47, 60, 50): 'ind', (60, 52, 60, 55): 'val', (60, 57, 60, 62): 'shape'}, {}), '(ind, val, shape)', True, 'from tensorflow.python.framework import sparse_tensor as sparse_tensor_lib\n'), ((71, 11, 71, 63), 'tensorflow.python.framework.sparse_tensor.SparseTensorValue', 'sparse_tensor_lib.SparseTensorValue', ({(71, 47, 71, 50): 'ind', (71, 52, 71, 55): 'val', (71, 57, 71, 62): 'shape'}, {}), '(ind, val, shape)', True, 'from tensorflow.python.framework import sparse_tensor as sparse_tensor_lib\n'), ((77, 11, 77, 63), 'tensorflow.python.framework.sparse_tensor.SparseTensorValue', 'sparse_tensor_lib.SparseTensorValue', ({(77, 47, 77, 50): 'ind', (77, 52, 77, 55): 'val', (77, 57, 77, 62): 'shape'}, {}), '(ind, val, shape)', True, 'from tensorflow.python.framework import sparse_tensor as sparse_tensor_lib\n'), ((187, 4, 187, 23), 'numpy.random.seed', 'np.random.seed', ({(187, 19, 187, 22): '(127)'}, {}), '(127)', True, 'import numpy as np\n'), ((190, 20, 191, 54), 'numpy.random.randint', 'np.random.randint', (), '', True, 'import numpy as np\n'), ((192, 20, 192, 59), 'numpy.arange', 'np.arange', (), '', True, 'import numpy as np\n'), ((196, 12, 196, 66), 'numpy.asarray', 'np.asarray', (), '', True, 'import numpy as np\n'), ((48, 8, 48, 43), 'tensorflow.python.ops.array_ops.placeholder', 'array_ops.placeholder', ({(48, 30, 48, 42): 'dtypes.int64'}, {}), '(dtypes.int64)', False, 'from tensorflow.python.ops import array_ops\n'), ((49, 8, 49, 36), 'tensorflow.python.ops.array_ops.placeholder', 'array_ops.placeholder', ({(49, 30, 49, 35): 'dtype'}, {}), '(dtype)', False, 'from tensorflow.python.ops import array_ops\n'), ((49, 38, 49, 73), 'tensorflow.python.ops.array_ops.placeholder', 'array_ops.placeholder', ({(49, 60, 49, 72): 'dtypes.int64'}, {}), '(dtypes.int64)', False, 'from tensorflow.python.ops import array_ops\n'), ((86, 23, 86, 58), 'tensorflow.python.ops.array_ops.stack', 'array_ops.stack', ({(86, 39, 86, 57): '[handle0, handle1]'}, {}), '([handle0, handle1])', False, 'from tensorflow.python.ops import array_ops\n'), ((111, 23, 112, 61), 'tensorflow.python.framework.ops.convert_to_tensor', 'ops.convert_to_tensor', (), '', False, 'from tensorflow.python.framework import ops\n'), ((131, 22, 131, 72), 'numpy.array', 'np.array', (), '', True, 'import numpy as np\n'), ((132, 21, 132, 49), 'numpy.array', 'np.array', ({(132, 30, 132, 48): "[b'a', b'b', b'c']"}, {}), "([b'a', b'b', b'c'])", True, 'import numpy as np\n'), ((133, 20, 133, 52), 'numpy.array', 'np.array', (), '', True, 'import numpy as np\n'), ((160, 22, 161, 61), 'tensorflow.python.framework.ops.convert_to_tensor', 'ops.convert_to_tensor', (), '', False, 'from tensorflow.python.framework import ops\n'), ((52, 10, 53, 28), 'numpy.array', 'np.array', ({(52, 19, 53, 27): '[[0, 0], [1, 0], [1, 3], [1, 4], [3, 2], [3, 3]]'}, {}), '([[0, 0], [1, 0], [1, 3], [1, 4], [3, 2], [3, 3]])', True, 'import numpy as np\n'), ((54, 10, 54, 43), 'numpy.array', 'np.array', ({(54, 19, 54, 42): '[0, 10, 13, 14, 32, 33]'}, {}), '([0, 10, 13, 14, 32, 33])', True, 'import numpy as np\n'), ((59, 12, 59, 28), 'numpy.array', 'np.array', ({(59, 21, 59, 27): '[5, 6]'}, {}), '([5, 6])', True, 'import numpy as np\n'), ((63, 10, 64, 28), 'numpy.array', 'np.array', ({(63, 19, 64, 27): '[[0, 0], [1, 0], [1, 2], [1, 3], [2, 2], [2, 3]]'}, {}), '([[0, 0], [1, 0], [1, 2], [1, 3], [2, 2], [2, 3]])', True, 'import numpy as np\n'), ((65, 10, 65, 43), 'numpy.array', 'np.array', ({(65, 19, 65, 42): '[0, 10, 13, 14, 32, 33]'}, {}), '([0, 10, 13, 14, 32, 33])', True, 'import numpy as np\n'), ((70, 12, 70, 28), 'numpy.array', 'np.array', ({(70, 21, 70, 27): '[3, 4]'}, {}), '([3, 4])', True, 'import numpy as np\n'), ((74, 10, 74, 31), 'numpy.array', 'np.array', ({(74, 19, 74, 30): '[[0, 0, 0]]'}, {}), '([[0, 0, 0]])', True, 'import numpy as np\n'), ((75, 10, 75, 23), 'numpy.array', 'np.array', ({(75, 19, 75, 22): '[0]'}, {}), '([0])', True, 'import numpy as np\n'), ((76, 12, 76, 31), 'numpy.array', 'np.array', ({(76, 21, 76, 30): '[3, 4, 5]'}, {}), '([3, 4, 5])', True, 'import numpy as np\n'), ((81, 46, 81, 58), 'numpy.arange', 'np.arange', ({(81, 56, 81, 57): '6'}, {}), '(6)', True, 'import numpy as np\n'), ((82, 46, 82, 58), 'numpy.arange', 'np.arange', ({(82, 56, 82, 57): '6'}, {}), '(6)', True, 'import numpy as np\n'), ((104, 47, 104, 59), 'numpy.arange', 'np.arange', ({(104, 57, 104, 58): '6'}, {}), '(6)', True, 'import numpy as np\n'), ((105, 47, 105, 59), 'numpy.arange', 'np.arange', ({(105, 57, 105, 58): '6'}, {}), '(6)', True, 'import numpy as np\n'), ((153, 47, 153, 59), 'numpy.arange', 'np.arange', ({(153, 57, 153, 58): '6'}, {}), '(6)', True, 'import numpy as np\n'), ((173, 46, 173, 58), 'numpy.arange', 'np.arange', ({(173, 56, 173, 57): '6'}, {}), '(6)', True, 'import numpy as np\n'), ((198, 11, 198, 31), 'tensorflow.python.framework.ops.device', 'ops.device', ({(198, 22, 198, 30): '"""/cpu:0"""'}, {}), "('/cpu:0')", False, 'from tensorflow.python.framework import ops\n'), ((199, 18, 199, 45), 'tensorflow.python.ops.variables.Variable', 'variables.Variable', ({(199, 37, 199, 44): 'indices'}, {}), '(indices)', False, 'from tensorflow.python.ops import variables\n'), ((200, 17, 200, 43), 'tensorflow.python.ops.variables.Variable', 'variables.Variable', ({(200, 36, 200, 42): 'values'}, {}), '(values)', False, 'from tensorflow.python.ops import variables\n'), ((201, 16, 201, 41), 'tensorflow.python.ops.variables.Variable', 'variables.Variable', ({(201, 35, 201, 40): 'shape'}, {}), '(shape)', False, 'from tensorflow.python.ops import variables\n'), ((202, 13, 202, 67), 'tensorflow.python.framework.sparse_tensor.SparseTensor', 'sparse_tensor_lib.SparseTensor', ({(202, 44, 202, 51): 'indices', (202, 53, 202, 59): 'values', (202, 61, 202, 66): 'shape'}, {}), '(indices, values, shape)', True, 'from tensorflow.python.framework import sparse_tensor as sparse_tensor_lib\n'), ((209, 24, 209, 60), 'tensorflow.python.ops.sparse_ops.serialize_many_sparse', 'sparse_ops.serialize_many_sparse', ({(209, 57, 209, 59): 'st'}, {}), '(st)', False, 'from tensorflow.python.ops import sparse_ops\n'), ((210, 26, 211, 46), 'tensorflow.python.ops.sparse_ops.deserialize_many_sparse', 'sparse_ops.deserialize_many_sparse', (), '', False, 'from tensorflow.python.ops import sparse_ops\n'), ((218, 8, 219, 62), 'numpy.testing.assert_equal', 'np.testing.assert_equal', ({(218, 32, 218, 58): 'st_roundtrip_values.values', (219, 32, 219, 61): 'st_deserialized_values.values'}, {}), '(st_roundtrip_values.values, st_deserialized_values.\n    values)', True, 'import numpy as np\n'), ((220, 8, 221, 63), 'numpy.testing.assert_equal', 'np.testing.assert_equal', ({(220, 32, 220, 59): 'st_roundtrip_values.indices', (221, 32, 221, 62): 'st_deserialized_values.indices'}, {}), '(st_roundtrip_values.indices, st_deserialized_values\n    .indices)', True, 'import numpy as np\n'), ((222, 8, 223, 67), 'numpy.testing.assert_equal', 'np.testing.assert_equal', ({(222, 32, 222, 63): 'st_roundtrip_values.dense_shape', (223, 32, 223, 66): 'st_deserialized_values.dense_shape'}, {}), '(st_roundtrip_values.dense_shape,\n    st_deserialized_values.dense_shape)', True, 'import numpy as np\n'), ((80, 28, 80, 39), 'tensorflow.python.framework.ops.Graph', 'ops.Graph', ({}, {}), '()', False, 'from tensorflow.python.framework import ops\n'), ((197, 32, 197, 60), 'tensorflow.python.platform.benchmark.benchmark_config', 'benchmark.benchmark_config', ({}, {}), '()', False, 'from tensorflow.python.platform import benchmark\n'), ((214, 8, 214, 48), 'tensorflow.python.ops.variables.global_variables_initializer', 'variables.global_variables_initializer', ({}, {}), '()', False, 'from tensorflow.python.ops import variables\n')]