tensorflow/compiler/tests/scatter_nd_op_test.py - third_party/github.com/tensorflow/tensorflow - Git at Google

 # 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 tensorflow.ops.tf.scatter_nd."""

 import functools

 import numpy as np

 from tensorflow.compiler.tests import xla_test
 from tensorflow.python.framework import errors
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
 from tensorflow.python.platform import test


 def _AsType(v, vtype):
   return v.astype(vtype) if isinstance(v, np.ndarray) else vtype(v)


 def _FlatInnerDims(tensor, ndims=2):
   shape = list(tensor.shape)
   return tensor.reshape(
       [functools.reduce(lambda x, y: x * y, shape[:-ndims + 1], 1)] +
       shape[-ndims + 1:])


 def _FlatOuterDims(tensor, ndims=2):
   shape = list(tensor.shape)
   return tensor.reshape(
       shape[:ndims - 1] +
       [functools.reduce(lambda x, y: x * y, shape[ndims - 1:], 1)])


 def _NumpyScatterNd(ref, indices, updates, op):
   ixdim = indices.shape[-1]
   num_updates = indices.size // ixdim
   total_nd = len(ref.shape)
   slice_size = 1
   for i in range(ixdim, total_nd):
     slice_size *= ref.shape[i]
   flat_indices = _FlatInnerDims(indices)
   flat_updates = updates.reshape((num_updates, slice_size))
   output_flat = _FlatOuterDims(ref, ixdim + 1)
   for ix_updates, ix_output in enumerate(flat_indices):
     ix_output = tuple(ix_output)
     output_flat[ix_output] = op(output_flat[ix_output],
                                 flat_updates[ix_updates])
   return output_flat.reshape(ref.shape)


 def _NumpyUpdate(indices, updates, shape):
   ref = np.zeros(shape, dtype=updates.dtype)
   return _NumpyScatterNd(ref, indices, updates, lambda p, u: u)


 class ScatterNdTest(xla_test.XLATestCase):

   def _VariableRankTest(self,
                         np_scatter,
                         tf_scatter,
                         vtype,
                         itype,
                         repeat_indices=False):
     np.random.seed(8)
     ref_shapes = [(3, 6), (3, 6), (3, 6, 9), (3, 6, 9), (3, 6, 9), (3, 6, 9)]
     indices_shapes = [(2,), (2, 2), (2,), (2, 2), (2, 3), (2, 3, 3)]
     for ref_shape, indices_shape in zip(ref_shapes, indices_shapes):
       num_updates = indices_shape[0]
       ixdim = indices_shape[-1]

       indexable_area_shape = ()
       for i in range(ixdim):
         indexable_area_shape += (ref_shape[i],)
       all_indices = [
           list(coord)
           for coord, _ in np.ndenumerate(np.empty(indexable_area_shape, vtype))
       ]
       np.random.shuffle(all_indices)
       indices = np.array(all_indices[:num_updates])

       if num_updates > 1 and repeat_indices:
         indices = indices[:num_updates // 2]
         for _ in range(num_updates - num_updates // 2):
           indices = np.append(
               indices, [indices[np.random.randint(num_updates // 2)]], axis=0)
         np.random.shuffle(indices)
       indices = _AsType(indices[:num_updates], itype)

       updates_shape = (num_updates,)
       for i in range(ixdim, len(ref_shape)):
         updates_shape += (ref_shape[i],)
       updates = _AsType(np.random.randn(*(updates_shape)), vtype)

       # Scatter via numpy
       np_out = np_scatter(indices, updates, ref_shape)
       # Scatter via tensorflow
       tf_out = tf_scatter(indices, updates, ref_shape)

       self.assertAllClose(np_out, tf_out)

   def _VariableRankTests(self, np_scatter, tf_scatter):
     for vtype in self.numeric_types:
       for itype in set([np.int32, np.int64]).intersection(set(self.int_types)):
         self._VariableRankTest(np_scatter, tf_scatter, vtype, itype)

   def _runScatterNd(self, indices, updates, shape):
     with self.session():
       updates_placeholder = array_ops.placeholder(updates.dtype)
       indices_placeholder = array_ops.placeholder(indices.dtype)
       with self.test_scope():
         output = array_ops.scatter_nd(indices_placeholder, updates_placeholder,
                                       shape)
       feed_dict = {updates_placeholder: updates, indices_placeholder: indices}
       return output.eval(feed_dict=feed_dict)

   def testSimple(self):
     indices = np.array([[4], [3], [1], [7]], dtype=np.int32)
     updates = np.array([9, 10, 11, 12], dtype=np.float32)
     expected = np.array([0, 11, 0, 10, 9, 0, 0, 12], dtype=np.int32)
     self.assertAllEqual(expected, self._runScatterNd(indices, updates, [8]))

   def testRepeatedIndices(self):
     indices = np.array([[0], [1], [0], [1]], dtype=np.int32)
     updates = np.array([9, 10, 11, 12], dtype=np.float32)
     expected = np.array([20, 22], dtype=np.int32)
     self.assertAllEqual(expected, self._runScatterNd(indices, updates, [2]))

   def testSimple2(self):
     indices = np.array([[1, 0], [1, 1]], dtype=np.int32)
     updates = np.array([11., 12.], dtype=np.float32)
     expected = np.array([[0., 0.], [11., 12.], [0., 0.]], dtype=np.float32)
     self.assertAllEqual(expected, self._runScatterNd(indices, updates, [3, 2]))

   def testSimple3(self):
     indices = np.array([[1]], dtype=np.int32)
     updates = np.array([[11., 12.]], dtype=np.float32)
     expected = np.array([[0., 0.], [11., 12.], [0., 0.]])
     self.assertAllEqual(expected, self._runScatterNd(indices, updates, [3, 2]))

   def testVariableRankUpdate(self):
     self._VariableRankTests(_NumpyUpdate, self._runScatterNd)

   def testExtraIndicesDimensions(self):
     indices = np.zeros([1, 1, 2], np.int32)
     updates = np.zeros([1, 1], np.int32)
     expected = np.zeros([2, 2], dtype=np.int32)
     self.assertAllEqual(expected, self._runScatterNd(indices, updates, [2, 2]))

   @test_util.disable_mlir_bridge("Error messages differ")
   def testRank3InvalidShape1(self):
     indices = np.zeros([3, 2, 2], np.int32)
     updates = np.zeros([2, 2, 2], np.int32)
     with self.assertRaisesWithPredicateMatch(errors.InvalidArgumentError,
                                              "Must have updates.shape"):
       self._runScatterNd(indices, updates, [2, 2, 2])

   @test_util.disable_mlir_bridge("Error messages differ")
   def testRank3InvalidShape2(self):
     indices = np.zeros([2, 2, 1], np.int32)
     updates = np.zeros([2, 2], np.int32)
     with self.assertRaisesWithPredicateMatch(errors.InvalidArgumentError,
                                              "Must have updates.shape"):
       self._runScatterNd(indices, updates, [2, 2, 2])

   def testScatterOutOfRange(self):
     updates = np.array([-3, -4, -5]).astype(np.float32)

     # Indices all in range, no problem.
     indices = np.array([[2], [0], [5]], dtype=np.int32)
     self._runScatterNd(indices, updates, [6])

     # Indices out of range should not fail. It produces implementation-defined
     # output.
     indices = np.array([[-1], [0], [5]], dtype=np.int32)
     self._runScatterNd(indices, updates, [6])
     indices = np.array([[2], [0], [6]], dtype=np.int32)
     self._runScatterNd(indices, updates, [6])


 class ScatterNdTensorTest(xla_test.XLATestCase):

   def _runScatter(self, op):
     indices_np = np.array([[4], [3], [1], [7]], dtype=np.int32)
     updates_np = np.array([9, 10, 11, 12], dtype=np.float32)
     with self.session() as sess, self.test_scope():
       indices = array_ops.placeholder(indices_np.dtype, shape=indices_np.shape)
       updates = array_ops.placeholder(updates_np.dtype, shape=updates_np.shape)
       t = array_ops.ones([8], dtype=np.float32)

       out = op(t, indices, updates)
       return sess.run(out, feed_dict={indices: indices_np, updates: updates_np})

   def testAdd(self):
     self.assertAllEqual(
         self._runScatter(array_ops.tensor_scatter_add),
         np.array([1, 12, 1, 11, 10, 1, 1, 13], dtype=np.float32))

   def testSub(self):
     self.assertAllEqual(
         self._runScatter(array_ops.tensor_scatter_sub),
         np.array([1, -10, 1, -9, -8, 1, 1, -11], dtype=np.float32))

   def testUpdate(self):
     self.assertAllEqual(
         self._runScatter(array_ops.tensor_scatter_update),
         np.array([1, 11, 1, 10, 9, 1, 1, 12], dtype=np.float32))


 class ScatterNdTensorScalarUpdateTest(xla_test.XLATestCase):

   def _runScatter(self, op):
     indices_np = np.array([[4], [3], [1], [7]], dtype=np.int32)
     updates_np = np.array(9, dtype=np.float32)
     with self.session() as sess, self.test_scope():
       indices = array_ops.placeholder(indices_np.dtype, shape=indices_np.shape)
       updates = array_ops.placeholder(updates_np.dtype, shape=updates_np.shape)
       t = array_ops.ones([8], dtype=np.float32)

       out = op(t, indices, updates)
       return sess.run(out, feed_dict={indices: indices_np, updates: updates_np})

   def testUpdate(self):
     self.assertAllEqual(
         self._runScatter(array_ops.tensor_scatter_update),
         np.array([1, 9, 1, 9, 9, 1, 1, 9], dtype=np.float32))

   def testAdd(self):
     self.assertAllEqual(
         self._runScatter(array_ops.tensor_scatter_add),
         np.array([1, 10, 1, 10, 10, 1, 1, 10], dtype=np.float32))

 if __name__ == "__main__":
   test.main()
	# 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 tensorflow.ops.tf.scatter_nd."""

	import functools

	import numpy as np

	from tensorflow.compiler.tests import xla_test
	from tensorflow.python.framework import errors
	from tensorflow.python.framework import test_util
	from tensorflow.python.ops import array_ops
	from tensorflow.python.platform import test


	def _AsType(v, vtype):
	return v.astype(vtype) if isinstance(v, np.ndarray) else vtype(v)


	def _FlatInnerDims(tensor, ndims=2):
	shape = list(tensor.shape)
	return tensor.reshape(
	[functools.reduce(lambda x, y: x * y, shape[:-ndims + 1], 1)] +
	shape[-ndims + 1:])


	def _FlatOuterDims(tensor, ndims=2):
	shape = list(tensor.shape)
	return tensor.reshape(
	shape[:ndims - 1] +
	[functools.reduce(lambda x, y: x * y, shape[ndims - 1:], 1)])


	def _NumpyScatterNd(ref, indices, updates, op):
	ixdim = indices.shape[-1]
	num_updates = indices.size // ixdim
	total_nd = len(ref.shape)
	slice_size = 1
	for i in range(ixdim, total_nd):
	slice_size *= ref.shape[i]
	flat_indices = _FlatInnerDims(indices)
	flat_updates = updates.reshape((num_updates, slice_size))
	output_flat = _FlatOuterDims(ref, ixdim + 1)
	for ix_updates, ix_output in enumerate(flat_indices):
	ix_output = tuple(ix_output)
	output_flat[ix_output] = op(output_flat[ix_output],
	flat_updates[ix_updates])
	return output_flat.reshape(ref.shape)


	def _NumpyUpdate(indices, updates, shape):
	ref = np.zeros(shape, dtype=updates.dtype)
	return _NumpyScatterNd(ref, indices, updates, lambda p, u: u)


	class ScatterNdTest(xla_test.XLATestCase):

	def _VariableRankTest(self,
	np_scatter,
	tf_scatter,
	vtype,
	itype,
	repeat_indices=False):
	np.random.seed(8)
	ref_shapes = [(3, 6), (3, 6), (3, 6, 9), (3, 6, 9), (3, 6, 9), (3, 6, 9)]
	indices_shapes = [(2,), (2, 2), (2,), (2, 2), (2, 3), (2, 3, 3)]
	for ref_shape, indices_shape in zip(ref_shapes, indices_shapes):
	num_updates = indices_shape[0]
	ixdim = indices_shape[-1]

	indexable_area_shape = ()
	for i in range(ixdim):
	indexable_area_shape += (ref_shape[i],)
	all_indices = [
	list(coord)
	for coord, _ in np.ndenumerate(np.empty(indexable_area_shape, vtype))
	]
	np.random.shuffle(all_indices)
	indices = np.array(all_indices[:num_updates])

	if num_updates > 1 and repeat_indices:
	indices = indices[:num_updates // 2]
	for _ in range(num_updates - num_updates // 2):
	indices = np.append(
	indices, [indices[np.random.randint(num_updates // 2)]], axis=0)
	np.random.shuffle(indices)
	indices = _AsType(indices[:num_updates], itype)

	updates_shape = (num_updates,)
	for i in range(ixdim, len(ref_shape)):
	updates_shape += (ref_shape[i],)
	updates = _AsType(np.random.randn(*(updates_shape)), vtype)

	# Scatter via numpy
	np_out = np_scatter(indices, updates, ref_shape)
	# Scatter via tensorflow
	tf_out = tf_scatter(indices, updates, ref_shape)

	self.assertAllClose(np_out, tf_out)

	def _VariableRankTests(self, np_scatter, tf_scatter):
	for vtype in self.numeric_types:
	for itype in set([np.int32, np.int64]).intersection(set(self.int_types)):
	self._VariableRankTest(np_scatter, tf_scatter, vtype, itype)

	def _runScatterNd(self, indices, updates, shape):
	with self.session():
	updates_placeholder = array_ops.placeholder(updates.dtype)
	indices_placeholder = array_ops.placeholder(indices.dtype)
	with self.test_scope():
	output = array_ops.scatter_nd(indices_placeholder, updates_placeholder,
	shape)
	feed_dict = {updates_placeholder: updates, indices_placeholder: indices}
	return output.eval(feed_dict=feed_dict)

	def testSimple(self):
	indices = np.array([[4], [3], [1], [7]], dtype=np.int32)
	updates = np.array([9, 10, 11, 12], dtype=np.float32)
	expected = np.array([0, 11, 0, 10, 9, 0, 0, 12], dtype=np.int32)
	self.assertAllEqual(expected, self._runScatterNd(indices, updates, [8]))

	def testRepeatedIndices(self):
	indices = np.array([[0], [1], [0], [1]], dtype=np.int32)
	updates = np.array([9, 10, 11, 12], dtype=np.float32)
	expected = np.array([20, 22], dtype=np.int32)
	self.assertAllEqual(expected, self._runScatterNd(indices, updates, [2]))

	def testSimple2(self):
	indices = np.array([[1, 0], [1, 1]], dtype=np.int32)
	updates = np.array([11., 12.], dtype=np.float32)
	expected = np.array([[0., 0.], [11., 12.], [0., 0.]], dtype=np.float32)
	self.assertAllEqual(expected, self._runScatterNd(indices, updates, [3, 2]))

	def testSimple3(self):
	indices = np.array([[1]], dtype=np.int32)
	updates = np.array([[11., 12.]], dtype=np.float32)
	expected = np.array([[0., 0.], [11., 12.], [0., 0.]])
	self.assertAllEqual(expected, self._runScatterNd(indices, updates, [3, 2]))

	def testVariableRankUpdate(self):
	self._VariableRankTests(_NumpyUpdate, self._runScatterNd)

	def testExtraIndicesDimensions(self):
	indices = np.zeros([1, 1, 2], np.int32)
	updates = np.zeros([1, 1], np.int32)
	expected = np.zeros([2, 2], dtype=np.int32)
	self.assertAllEqual(expected, self._runScatterNd(indices, updates, [2, 2]))

	@test_util.disable_mlir_bridge("Error messages differ")
	def testRank3InvalidShape1(self):
	indices = np.zeros([3, 2, 2], np.int32)
	updates = np.zeros([2, 2, 2], np.int32)
	with self.assertRaisesWithPredicateMatch(errors.InvalidArgumentError,
	"Must have updates.shape"):
	self._runScatterNd(indices, updates, [2, 2, 2])

	@test_util.disable_mlir_bridge("Error messages differ")
	def testRank3InvalidShape2(self):
	indices = np.zeros([2, 2, 1], np.int32)
	updates = np.zeros([2, 2], np.int32)
	with self.assertRaisesWithPredicateMatch(errors.InvalidArgumentError,
	"Must have updates.shape"):
	self._runScatterNd(indices, updates, [2, 2, 2])

	def testScatterOutOfRange(self):
	updates = np.array([-3, -4, -5]).astype(np.float32)

	# Indices all in range, no problem.
	indices = np.array([[2], [0], [5]], dtype=np.int32)
	self._runScatterNd(indices, updates, [6])

	# Indices out of range should not fail. It produces implementation-defined
	# output.
	indices = np.array([[-1], [0], [5]], dtype=np.int32)
	self._runScatterNd(indices, updates, [6])
	indices = np.array([[2], [0], [6]], dtype=np.int32)
	self._runScatterNd(indices, updates, [6])


	class ScatterNdTensorTest(xla_test.XLATestCase):

	def _runScatter(self, op):
	indices_np = np.array([[4], [3], [1], [7]], dtype=np.int32)
	updates_np = np.array([9, 10, 11, 12], dtype=np.float32)
	with self.session() as sess, self.test_scope():
	indices = array_ops.placeholder(indices_np.dtype, shape=indices_np.shape)
	updates = array_ops.placeholder(updates_np.dtype, shape=updates_np.shape)
	t = array_ops.ones([8], dtype=np.float32)

	out = op(t, indices, updates)
	return sess.run(out, feed_dict={indices: indices_np, updates: updates_np})

	def testAdd(self):
	self.assertAllEqual(
	self._runScatter(array_ops.tensor_scatter_add),
	np.array([1, 12, 1, 11, 10, 1, 1, 13], dtype=np.float32))

	def testSub(self):
	self.assertAllEqual(
	self._runScatter(array_ops.tensor_scatter_sub),
	np.array([1, -10, 1, -9, -8, 1, 1, -11], dtype=np.float32))

	def testUpdate(self):
	self.assertAllEqual(
	self._runScatter(array_ops.tensor_scatter_update),
	np.array([1, 11, 1, 10, 9, 1, 1, 12], dtype=np.float32))


	class ScatterNdTensorScalarUpdateTest(xla_test.XLATestCase):

	def _runScatter(self, op):
	indices_np = np.array([[4], [3], [1], [7]], dtype=np.int32)
	updates_np = np.array(9, dtype=np.float32)
	with self.session() as sess, self.test_scope():
	indices = array_ops.placeholder(indices_np.dtype, shape=indices_np.shape)
	updates = array_ops.placeholder(updates_np.dtype, shape=updates_np.shape)
	t = array_ops.ones([8], dtype=np.float32)

	out = op(t, indices, updates)
	return sess.run(out, feed_dict={indices: indices_np, updates: updates_np})

	def testUpdate(self):
	self.assertAllEqual(
	self._runScatter(array_ops.tensor_scatter_update),
	np.array([1, 9, 1, 9, 9, 1, 1, 9], dtype=np.float32))

	def testAdd(self):
	self.assertAllEqual(
	self._runScatter(array_ops.tensor_scatter_add),
	np.array([1, 10, 1, 10, 10, 1, 1, 10], dtype=np.float32))

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