1# Copyright 2015 The TensorFlow Authors. All Rights Reserved. 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14# ============================================================================== 15"""Functional tests for segment reduction ops.""" 16 17from __future__ import absolute_import 18from __future__ import division 19from __future__ import print_function 20 21import itertools 22 23import numpy as np 24 25from tensorflow.python.client import session 26from tensorflow.python.framework import constant_op 27from tensorflow.python.framework import dtypes as dtypes_lib 28from tensorflow.python.framework import errors_impl 29from tensorflow.python.framework import ops 30from tensorflow.python.framework import test_util 31from tensorflow.python.ops import gradient_checker 32from tensorflow.python.ops import gradient_checker_v2 33from tensorflow.python.ops import math_ops 34from tensorflow.python.ops import variables 35from tensorflow.python.platform import test 36 37 38class SegmentReductionHelper(test.TestCase): 39 40 def _input(self, input_shape, dtype=dtypes_lib.int32): 41 num_elem = 1 42 for x in input_shape: 43 num_elem *= x 44 values = np.arange(1, num_elem + 1) 45 np_values = values.reshape(input_shape).astype(dtype.as_numpy_dtype) 46 # Add a non-zero imaginary component to complex types. 47 if dtype.is_complex: 48 np_values -= 1j * np_values 49 return constant_op.constant( 50 np_values, shape=input_shape, dtype=dtype), np_values 51 52 def _segmentReduce(self, indices, x, op1, op2=None, num_segments=None, 53 initial_value=0): 54 if not x.size: 55 return np.array([]) 56 indices = np.asarray(indices) 57 if num_segments is None: 58 num_segments = indices[-1] + 1 59 output = [None] * num_segments 60 slice_shape = x.shape[indices.ndim:] 61 x_flat = x.reshape((indices.size,) + slice_shape) 62 for i, index in enumerate(indices.ravel()): 63 if (output[index] is not None) and op1 == np.max: 64 for j in range(0, output[index].shape[0]): 65 output[index][j] = op1([output[index][j], x_flat[i][j]]) 66 elif output[index] is not None: 67 output[index] = op1(output[index], x_flat[i]) 68 else: 69 output[index] = x_flat[i] 70 # zero initialize values that are still uncalculated. 71 initial_value_slice = np.ones(slice_shape) * initial_value 72 output = [o if o is not None else initial_value_slice for o in output] 73 if op2 is not None: 74 output = [op2(o) for o in output] 75 output = [o.reshape(slice_shape) for o in output] 76 return np.array(output) 77 78 def _mean_cum_op(self, x, y): 79 return (x[0] + y, x[1] + 1) if isinstance(x, tuple) else (x + y, 2) 80 81 def _mean_reduce_op(self, x): 82 return x[0] / x[1] if isinstance(x, tuple) else x 83 84 def _sqrt_n_reduce_op(self, x): 85 return x[0] / np.sqrt(x[1]) if isinstance(x, tuple) else x 86 87 88class SegmentReductionOpTest(SegmentReductionHelper): 89 90 def testValues(self): 91 dtypes = [ 92 dtypes_lib.float32, dtypes_lib.float64, dtypes_lib.int64, 93 dtypes_lib.int32, dtypes_lib.complex64, dtypes_lib.complex128 94 ] 95 96 # Each item is np_op1, np_op2, tf_op 97 ops_list = [(np.add, None, math_ops.segment_sum), 98 (self._mean_cum_op, self._mean_reduce_op, 99 math_ops.segment_mean), 100 (np.ndarray.__mul__, None, math_ops.segment_prod), 101 (np.minimum, None, math_ops.segment_min), 102 (np.maximum, None, math_ops.segment_max)] 103 104 # A subset of ops has been enabled for complex numbers 105 complex_ops_list = [(np.add, None, math_ops.segment_sum), 106 (np.ndarray.__mul__, None, math_ops.segment_prod), 107 (self._mean_cum_op, self._mean_reduce_op, 108 math_ops.segment_mean)] 109 110 n = 10 111 shape = [n, 2] 112 indices = [i // 3 for i in range(n)] 113 for dtype in dtypes: 114 if dtype in (dtypes_lib.complex64, dtypes_lib.complex128): 115 curr_ops_list = complex_ops_list 116 else: 117 curr_ops_list = ops_list 118 for use_gpu in [True, False]: 119 with self.cached_session(use_gpu=use_gpu): 120 tf_x, np_x = self._input(shape, dtype=dtype) 121 for np_op1, np_op2, tf_op in curr_ops_list: 122 np_ans = self._segmentReduce(indices, np_x, np_op1, np_op2) 123 s = tf_op(data=tf_x, segment_ids=indices) 124 tf_ans = self.evaluate(s) 125 self.assertAllClose(np_ans, tf_ans) 126 # NOTE(mrry): The static shape inference that computes 127 # `tf_ans.shape` can only infer that sizes from dimension 1 128 # onwards, because the size of dimension 0 is data-dependent 129 # and may therefore vary dynamically. 130 self.assertAllEqual(np_ans.shape[1:], tf_ans.shape[1:]) 131 132 @test_util.run_deprecated_v1 133 def testSegmentIdsShape(self): 134 shape = [4, 4] 135 tf_x, _ = self._input(shape) 136 indices = constant_op.constant([0, 1, 2, 2], shape=[2, 2]) 137 with self.assertRaises(ValueError): 138 math_ops.segment_sum(data=tf_x, segment_ids=indices) 139 140 @test_util.run_deprecated_v1 141 def testSegmentIdsSize(self): 142 shape = [4, 4] 143 for use_gpu in [True, False]: 144 with self.cached_session(use_gpu=use_gpu): 145 tf_x, _ = self._input(shape) 146 indices = [0, 1] 147 s = math_ops.segment_sum(data=tf_x, segment_ids=indices) 148 with self.assertRaisesOpError("segment_ids should be the same size"): 149 self.evaluate(s) 150 151 @test_util.run_deprecated_v1 152 def testSegmentIdsValid(self): 153 # This is a baseline for the following SegmentIdsInvalid* tests. 154 shape = [4, 4] 155 for use_gpu in [True, False]: 156 with self.cached_session(use_gpu=use_gpu): 157 tf_x, _ = self._input(shape, dtype=dtypes_lib.float32) 158 indices = [0, 0, 0, 1] 159 result = math_ops.segment_sum(data=tf_x, segment_ids=indices).eval() 160 self.assertAllEqual([[15, 18, 21, 24], [13, 14, 15, 16]], result) 161 162 def testSegmentIdsGreaterThanZero(self): 163 shape = [4, 4] 164 for use_gpu in [True, False]: 165 with self.cached_session(use_gpu=use_gpu): 166 tf_x, np_x = self._input(shape, dtype=dtypes_lib.float32) 167 indices = [1, 1, 2, 2] 168 np_ans = self._segmentReduce(indices, np_x, np.add) 169 s = math_ops.segment_sum(data=tf_x, segment_ids=indices) 170 tf_ans = self.evaluate(s) 171 self.assertAllClose(np_ans, tf_ans) 172 173 def testSegmentIdsHole(self): 174 shape = [4, 4] 175 for use_gpu in [True, False]: 176 with self.cached_session(use_gpu=use_gpu): 177 tf_x, np_x = self._input(shape, dtype=dtypes_lib.float32) 178 indices = [0, 0, 3, 3] 179 np_ans = self._segmentReduce(indices, np_x, np.add) 180 s = math_ops.segment_sum(data=tf_x, segment_ids=indices) 181 tf_ans = self.evaluate(s) 182 self.assertAllClose(np_ans, tf_ans) 183 184 @test_util.run_deprecated_v1 185 def testSegmentIdsInvalid1(self): 186 shape = [4, 4] 187 with self.cached_session(): 188 tf_x, _ = self._input(shape) 189 indices = [-1, -1, 0, 0] 190 s = math_ops.segment_sum(data=tf_x, segment_ids=indices) 191 with self.assertRaisesOpError( 192 r"Segment id -1 out of range \[0, 1\), possibly because " 193 "'segment_ids' input is not sorted."): 194 self.evaluate(s) 195 196 @test_util.run_deprecated_v1 197 def testSegmentIdsInvalid2(self): 198 shape = [4, 4] 199 with self.cached_session(): 200 tf_x, _ = self._input(shape) 201 indices = [0, 1, 0, 1] 202 s = math_ops.segment_sum(data=tf_x, segment_ids=indices) 203 with self.assertRaisesOpError("segment ids are not increasing"): 204 self.evaluate(s) 205 206 @test_util.run_deprecated_v1 207 def testSegmentIdsInvalid3(self): 208 shape = [4, 4] 209 with self.cached_session(): 210 tf_x, _ = self._input(shape) 211 indices = [0, 1, 2, 0] 212 s = math_ops.segment_sum(data=tf_x, segment_ids=indices) 213 with self.assertRaisesOpError( 214 r"Segment id 1 out of range \[0, 1\), possibly " 215 "because 'segment_ids' input is not sorted."): 216 self.evaluate(s) 217 218 @test_util.run_deprecated_v1 219 def testSegmentIdsInvalid4(self): 220 shape = [4, 4] 221 for use_gpu in [True, False]: 222 with self.cached_session(use_gpu=use_gpu): 223 tf_x, _ = self._input(shape, dtype=dtypes_lib.float32) 224 indices = [0, 0, 0, -1] 225 s = math_ops.segment_sum(data=tf_x, segment_ids=indices) 226 with self.assertRaisesOpError("segment ids must be >= 0"): 227 self.evaluate(s) 228 229 @test_util.run_deprecated_v1 230 def testSegmentIdsInvalid5(self): 231 shape = [4, 4] 232 for use_gpu in [True, False]: 233 with self.cached_session(use_gpu=use_gpu): 234 tf_x, _ = self._input(shape, dtype=dtypes_lib.float32) 235 indices = [0, 0, 0, -2] 236 s = math_ops.segment_sum(data=tf_x, segment_ids=indices) 237 with self.assertRaisesOpError("segment ids must be >= 0"): 238 self.evaluate(s) 239 240 @test_util.run_deprecated_v1 241 def testGradient(self): 242 shape = [4, 4] 243 indices = [0, 1, 2, 2] 244 for tf_op in [ 245 math_ops.segment_sum, math_ops.segment_mean, math_ops.segment_min, 246 math_ops.segment_max 247 ]: 248 with self.cached_session(): 249 tf_x, np_x = self._input(shape, dtype=dtypes_lib.float64) 250 s = tf_op(data=tf_x, segment_ids=indices) 251 jacob_t, jacob_n = gradient_checker.compute_gradient( 252 tf_x, 253 shape, 254 s, [3, 4], 255 x_init_value=np_x.astype(np.double), 256 delta=1) 257 self.assertAllClose(jacob_t, jacob_n) 258 259 def testDataInvalid(self): 260 # Test case for GitHub issue 40653. 261 for use_gpu in [True, False]: 262 with self.cached_session(use_gpu=use_gpu): 263 with self.assertRaisesRegex( 264 (ValueError, errors_impl.InvalidArgumentError), 265 "must be at least rank 1"): 266 s = math_ops.segment_mean( 267 data=np.uint16(10), segment_ids=np.array([]).astype("int64")) 268 self.evaluate(s) 269 270 271class UnsortedSegmentTest(SegmentReductionHelper): 272 273 def __init__(self, methodName='runTest'): 274 # Each item is np_op1, np_op2, tf_op, initial_value functor 275 self.ops_list = [(np.add, None, 276 math_ops.unsorted_segment_sum, lambda t: 0), 277 (self._mean_cum_op, self._mean_reduce_op, 278 math_ops.unsorted_segment_mean, lambda t: 0), 279 (self._mean_cum_op, self._sqrt_n_reduce_op, 280 math_ops.unsorted_segment_sqrt_n, lambda t: 0), 281 (np.ndarray.__mul__, None, 282 math_ops.unsorted_segment_prod, lambda t: 1), 283 (np.minimum, None, 284 math_ops.unsorted_segment_min, lambda t: t.max), 285 (np.maximum, None, 286 math_ops.unsorted_segment_max, lambda t: t.min)] 287 288 # A subset of ops has been enabled for complex numbers 289 self.complex_ops_list = [(np.add, None, 290 math_ops.unsorted_segment_sum, lambda t: 0), 291 (np.ndarray.__mul__, None, 292 math_ops.unsorted_segment_prod, lambda t: 1)] 293 self.differentiable_dtypes = [dtypes_lib.float16, dtypes_lib.float32, 294 dtypes_lib.float64] 295 self.all_dtypes = (self.differentiable_dtypes + 296 [dtypes_lib.bfloat16, 297 dtypes_lib.int64, dtypes_lib.int32, 298 dtypes_lib.complex64, dtypes_lib.complex128]) 299 super(UnsortedSegmentTest, self).__init__(methodName=methodName) 300 301 def testValues(self): 302 indices_flat = np.array([0, 4, 0, 8, 3, 8, 4, 7, 7, 3]) 303 num_segments = 12 304 for indices in indices_flat, indices_flat.reshape(5, 2): 305 shape = indices.shape + (2,) 306 for dtype in self.all_dtypes: 307 ops_list = self.complex_ops_list if dtype.is_complex else self.ops_list 308 tf_x, np_x = self._input(shape, dtype=dtype) 309 for use_gpu in [True, False]: 310 with self.cached_session(): 311 for np_op1, np_op2, tf_op, init_op in ops_list: 312 # sqrt_n doesn't support integers 313 if (np_op2 == self._sqrt_n_reduce_op and dtype.is_integer): 314 continue 315 # todo(philjd): enable this test once real_div supports bfloat16 316 if (np_op2 in [self._sqrt_n_reduce_op, self._mean_reduce_op] and 317 dtype == dtypes_lib.bfloat16): 318 continue 319 np_ans = self._segmentReduce( 320 indices, np_x, np_op1, np_op2, num_segments=num_segments, 321 initial_value=init_op(dtype)) 322 s = tf_op(tf_x, segment_ids=indices, num_segments=num_segments) 323 tf_ans = self.evaluate(s) 324 if dtype is dtypes_lib.bfloat16: 325 tf_ans = tf_ans.astype(np.float32) 326 self.assertAllCloseAccordingToType(np_ans, tf_ans) 327 self.assertShapeEqual(np_ans, s) 328 329 def testNumSegmentsTypes(self): 330 dtypes = [dtypes_lib.int32, dtypes_lib.int64] 331 indices_flat = np.array([0, 4, 0, 8, 3, 8, 4, 7, 7, 3]) 332 num_segments = 12 333 for indices in indices_flat, indices_flat.reshape(5, 2): 334 shape = indices.shape + (2,) 335 for dtype in dtypes: 336 with self.cached_session(): 337 tf_x, np_x = self._input(shape) 338 num_segments_constant = constant_op.constant( 339 num_segments, dtype=dtype) 340 np_ans = self._segmentReduce( 341 indices, np_x, np.add, op2=None, num_segments=num_segments) 342 s = math_ops.unsorted_segment_sum( 343 data=tf_x, 344 segment_ids=indices, 345 num_segments=num_segments_constant) 346 tf_ans = self.evaluate(s) 347 self.assertAllClose(np_ans, tf_ans) 348 self.assertShapeEqual(np_ans, s) 349 350 @test_util.run_deprecated_v1 351 def testGradientsTFGradients(self): 352 num_cols = 2 353 indices_flat = np.array([0, 4, 0, -1, 3, -1, 4, 7, 7, 3]) 354 num_segments = max(indices_flat) + 3 355 for dtype in self.differentiable_dtypes: 356 ops_list = self.complex_ops_list if dtype.is_complex else self.ops_list 357 for indices in indices_flat, indices_flat.reshape(5, 2): 358 shape = indices.shape + (num_cols,) 359 # test CPU and GPU as tf.gather behaves differently on each device 360 for use_gpu in [False, True]: 361 with self.cached_session(use_gpu=use_gpu): 362 for _, _, tf_op, _ in ops_list: 363 tf_x, np_x = self._input(shape, dtype=dtype) 364 s = tf_op(tf_x, indices, num_segments) 365 jacob_t, jacob_n = gradient_checker.compute_gradient( 366 tf_x, 367 shape, 368 s, [num_segments, num_cols], 369 x_init_value=np_x, 370 delta=1.) 371 self.assertAllCloseAccordingToType(jacob_t, jacob_n, 372 half_atol=1e-2) 373 374 @test_util.run_in_graph_and_eager_modes 375 def testGradientsGradientTape(self): 376 num_cols = 2 377 indices_flat = np.array([0, 4, 0, -1, 3, -1, 4, 7, 7, 3]) 378 num_segments = max(indices_flat) + 3 379 for dtype in self.differentiable_dtypes: 380 ops_list = self.complex_ops_list if dtype.is_complex else self.ops_list 381 for indices in indices_flat, indices_flat.reshape(5, 2): 382 shape = indices.shape + (num_cols,) 383 # test CPU and GPU as tf.gather behaves differently on each device 384 for use_gpu in [test_util.use_gpu, test_util.force_cpu]: 385 with use_gpu(): 386 for _, _, tf_op, _ in ops_list: 387 _, np_x = self._input(shape, dtype=dtype) 388 # pylint: disable=cell-var-from-loop 389 def f(x): 390 return tf_op(x, indices, num_segments) 391 gradient_tape_jacob_t, jacob_n = ( 392 gradient_checker_v2.compute_gradient( 393 f, [np_x], delta=1.)) 394 # pylint: enable=cell-var-from-loop 395 self.assertAllCloseAccordingToType(jacob_n, gradient_tape_jacob_t, 396 half_atol=1e-2) 397 398 @test_util.run_deprecated_v1 399 def testProdGrad(self): 400 # additional test for the prod gradient to ensure correct handling of zeros 401 values = np.array([0, 0, 1, 0, 2, 2, 3, 3, 3], dtype=np.float32) 402 indices = np.array([0, 0, 0, 1, 1, 1, 2, 2, 2], dtype=np.int32) 403 indices_neg = np.array([-1, 0, 0, -1, 1, 1, -1, 2, 2], dtype=np.int32) 404 values_tf = constant_op.constant(values) 405 # ground truth partial derivatives 406 gradients_indices = np.zeros((9, 3), dtype=np.float32) 407 gradients_indices_neg = np.zeros((9, 3), dtype=np.float32) 408 # the derivative w.r.t. to the other segments is zero, so here we only 409 # explicitly set the grad values for the corresponding segment 410 gradients_indices[range(9), indices] = [0, 0, 0, 4, 0, 0, 9, 9, 9] 411 gradients_indices_neg[range(9), indices_neg] = [0, 1, 0, 0, 2, 2, 0, 3, 3] 412 for use_gpu in [False, True]: 413 with self.cached_session(use_gpu=use_gpu): 414 for ind, grad_gt in [(indices, gradients_indices), 415 (indices_neg, gradients_indices_neg)]: 416 s = math_ops.unsorted_segment_prod(values_tf, 417 constant_op.constant(ind), 3) 418 jacob_t, jacob_n = gradient_checker.compute_gradient( 419 values_tf, (9,), s, (3,), x_init_value=values, delta=1) 420 self.assertAllClose(jacob_t, jacob_n) 421 self.assertAllClose(jacob_t, grad_gt) 422 423 @test_util.run_deprecated_v1 424 def testGradientMatchesSegmentSum(self): 425 # Strategy: compute the gradient for UnsortedSegmentSum and SegmentSum 426 # and compare the outputs, which should be identical. 427 # NB: for this test to work, indices must be valid for SegmentSum, namely 428 # it must be sorted, the indices must be contiguous, and num_segments 429 # must be max(indices) + 1. 430 indices = [0, 0, 1, 1, 1, 2, 3, 4, 5] 431 n = len(indices) 432 num_cols = 2 433 shape = [n, num_cols] 434 num_segments = max(indices) + 1 435 for dtype in self.differentiable_dtypes: 436 with self.cached_session(): 437 tf_x, np_x = self._input(shape, dtype=dtype) 438 # Results from UnsortedSegmentSum 439 unsorted_s = math_ops.unsorted_segment_sum( 440 data=tf_x, segment_ids=indices, num_segments=num_segments) 441 unsorted_jacob_t, unsorted_jacob_n = ( 442 gradient_checker.compute_gradient(tf_x, shape, unsorted_s, 443 [num_segments, num_cols], 444 x_init_value=np_x, delta=1)) 445 446 # Results from SegmentSum 447 sorted_s = math_ops.segment_sum(data=tf_x, segment_ids=indices) 448 sorted_jacob_t, sorted_jacob_n = gradient_checker.compute_gradient( 449 tf_x, 450 shape, 451 sorted_s, [num_segments, num_cols], 452 x_init_value=np_x, 453 delta=1) 454 self.assertAllClose(unsorted_jacob_t, sorted_jacob_t) 455 self.assertAllClose(unsorted_jacob_n, sorted_jacob_n) 456 457 @test_util.run_deprecated_v1 458 def testBadIndices(self): 459 # Note: GPU kernel does not return the out-of-range error needed for this 460 # test, so this test is marked as cpu-only. 461 # Note: With PR #13055 a negative index will be ignored silently. 462 with self.session(use_gpu=False): 463 for bad in [[2]], [[7]]: 464 unsorted = math_ops.unsorted_segment_sum([[17]], bad, num_segments=2) 465 with self.assertRaisesOpError( 466 r"segment_ids\[0,0\] = %d is out of range \[0, 2\)" % bad[0][0]): 467 self.evaluate(unsorted) 468 469 @test_util.run_deprecated_v1 470 def testEmptySecondDimension(self): 471 dtypes = [np.float16, np.float32, np.float64, np.int64, np.int32, 472 np.complex64, np.complex128] 473 with self.session(): 474 for dtype in dtypes: 475 for itype in (np.int32, np.int64): 476 data = np.zeros((2, 0), dtype=dtype) 477 segment_ids = np.array([0, 1], dtype=itype) 478 unsorted = math_ops.unsorted_segment_sum(data, segment_ids, 2) 479 self.assertAllEqual(unsorted, np.zeros((2, 0), dtype=dtype)) 480 481 def testDropNegatives(self): 482 # Note: the test is done by replacing segment_ids with 8 to -1 483 # for index and replace values generated by numpy with 0. 484 indices_flat = np.array([0, 4, 0, 8, 3, 8, 4, 7, 7, 3]) 485 num_segments = 12 486 for indices in indices_flat, indices_flat.reshape(5, 2): 487 shape = indices.shape + (2,) 488 for dtype in self.all_dtypes: 489 with self.session(): 490 tf_x, np_x = self._input(shape, dtype=dtype) 491 np_ans = self._segmentReduce( 492 indices, np_x, np.add, op2=None, num_segments=num_segments) 493 # Replace np_ans[8] with 0 for the value 494 np_ans[8:] = 0 495 # Replace 8 with -1 in indices 496 np.place(indices, indices == 8, [-1]) 497 s = math_ops.unsorted_segment_sum( 498 data=tf_x, segment_ids=indices, num_segments=num_segments) 499 tf_ans = self.evaluate(s) 500 self.assertAllClose(np_ans, tf_ans) 501 self.assertShapeEqual(np_ans, s) 502 503 504class SparseSegmentReductionHelper(SegmentReductionHelper): 505 506 def _sparse_input(self, input_shape, num_indices, dtype=dtypes_lib.int32): 507 a, b = super(SparseSegmentReductionHelper, self)._input(input_shape, dtype) 508 indices = np.random.randint(0, input_shape[0], num_indices).astype(np.int32) 509 return (constant_op.constant( 510 indices, dtype=dtypes_lib.int32), indices, a, b) 511 512 def _sparseSegmentReduce(self, 513 x, 514 indices, 515 segment_indices, 516 op1, 517 op2=None, 518 num_segments=None): 519 return self._segmentReduce( 520 segment_indices, x[indices], op1, op2, num_segments=num_segments) 521 522 def _sparseSegmentReduceGrad(self, ygrad, indices, segment_ids, output_dim0, 523 mode): 524 assert mode in ("sum", "mean", "sqrtn") 525 if mode != "sum": 526 weights = np.zeros(ygrad.shape[0], ygrad.dtype) 527 for segment in segment_ids: 528 weights[segment] += 1 529 weights = 1. / weights if mode == "mean" else 1. / np.sqrt(weights) 530 xgrad = np.zeros([output_dim0, ygrad.shape[1]], ygrad.dtype) 531 for segment, index in zip(segment_ids, indices): 532 if mode == "sum": 533 xgrad[index] += ygrad[segment] 534 else: 535 xgrad[index] += ygrad[segment] * weights[segment] 536 return xgrad 537 538 539class SparseSegmentReductionOpTest(SparseSegmentReductionHelper): 540 541 def testValues(self): 542 dtypes = [ 543 dtypes_lib.float32, dtypes_lib.float64, dtypes_lib.int64, 544 dtypes_lib.int32 545 ] 546 547 index_dtypes = [dtypes_lib.int32, dtypes_lib.int64] 548 segment_ids_dtypes = [dtypes_lib.int32, dtypes_lib.int64] 549 550 mean_dtypes = [dtypes_lib.float32, dtypes_lib.float64] 551 552 # Each item is np_op1, np_op2, tf_op 553 ops_list = [(np.add, None, math_ops.sparse_segment_sum), 554 (self._mean_cum_op, self._mean_reduce_op, 555 math_ops.sparse_segment_mean)] 556 557 n = 400 558 # Note that the GPU implem has different paths for different inner sizes. 559 for inner_size in [1, 2, 3, 32]: 560 shape = [n, inner_size] 561 segment_indices = [] 562 for i in range(20): 563 for _ in range(i + 1): 564 segment_indices.append(i) 565 num_indices = len(segment_indices) 566 for dtype in dtypes: 567 for index_dtype in index_dtypes: 568 for segment_ids_dtype in segment_ids_dtypes: 569 with self.cached_session(): 570 tf_indices, np_indices, tf_x, np_x = self._sparse_input( 571 shape, num_indices, dtype=dtype) 572 for np_op1, np_op2, tf_op in ops_list: 573 if (tf_op == math_ops.sparse_segment_mean and 574 dtype not in mean_dtypes): 575 continue 576 np_ans = self._sparseSegmentReduce(np_x, np_indices, 577 segment_indices, np_op1, 578 np_op2) 579 s = tf_op( 580 data=tf_x, 581 indices=math_ops.cast(tf_indices, index_dtype), 582 segment_ids=math_ops.cast(segment_indices, 583 segment_ids_dtype)) 584 tf_ans = self.evaluate(s) 585 self.assertAllClose(np_ans, tf_ans) 586 # NOTE(mrry): The static shape inference that computes 587 # `tf_ans.shape` can only infer that sizes from dimension 1 588 # onwards, because the size of dimension 0 is data-dependent 589 # and may therefore vary dynamically. 590 self.assertAllEqual(np_ans.shape[1:], tf_ans.shape[1:]) 591 592 def testSegmentIdsHole(self): 593 tf_x, np_x = self._input([10, 4], dtype=dtypes_lib.float32) 594 ops_list = [(np.add, None, math_ops.sparse_segment_sum), ( 595 self._mean_cum_op, self._mean_reduce_op, math_ops.sparse_segment_mean)] 596 segment_indices = [0, 2, 2, 2] 597 tf_indices = [8, 3, 0, 9] 598 with self.session(): 599 for np_op1, np_op2, tf_op in ops_list: 600 np_ans = self._sparseSegmentReduce(np_x, tf_indices, segment_indices, 601 np_op1, np_op2) 602 s = tf_op(data=tf_x, indices=tf_indices, segment_ids=segment_indices) 603 tf_ans = self.evaluate(s) 604 self.assertAllClose(np_ans, tf_ans) 605 606 def testWithNumSegments(self): 607 tf_x, np_x = self._input([10, 4], dtype=dtypes_lib.float32) 608 ops_list = [(np.add, None, math_ops.sparse_segment_sum_with_num_segments), 609 (self._mean_cum_op, self._mean_reduce_op, 610 math_ops.sparse_segment_mean_with_num_segments)] 611 segment_indices = [0, 2, 2, 2] 612 tf_indices = [8, 3, 0, 9] 613 num_segments = 5 614 with self.session(): 615 for np_op1, np_op2, tf_op in ops_list: 616 np_ans = self._sparseSegmentReduce( 617 np_x, 618 tf_indices, 619 segment_indices, 620 np_op1, 621 np_op2, 622 num_segments=num_segments) 623 s = tf_op( 624 data=tf_x, 625 indices=tf_indices, 626 segment_ids=segment_indices, 627 num_segments=num_segments) 628 tf_ans = self.evaluate(s) 629 self.assertAllClose(np_ans, tf_ans) 630 631 def testWithEmptySegments(self): 632 tf_x = constant_op.constant([], shape=[0, 4], dtype=dtypes_lib.float32) 633 ops_list = [ 634 math_ops.sparse_segment_sum_with_num_segments, 635 math_ops.sparse_segment_mean_with_num_segments 636 ] 637 segment_indices = [] 638 tf_indices = [] 639 num_segments = 5 640 with self.session(): 641 for tf_op in ops_list: 642 s = tf_op( 643 data=tf_x, 644 indices=tf_indices, 645 segment_ids=segment_indices, 646 num_segments=num_segments) 647 tf_ans = self.evaluate(s) 648 self.assertAllClose(np.zeros([5, 4]), tf_ans) 649 650 @test_util.run_in_graph_and_eager_modes 651 def testSegmentScalarIdiRaisesInvalidArgumentError(self): 652 """Test for github #46897.""" 653 ops_list = [ 654 math_ops.sparse_segment_sum, 655 math_ops.sparse_segment_mean, 656 math_ops.sparse_segment_sqrt_n, 657 ] 658 for op in ops_list: 659 with self.assertRaisesRegex( 660 (ValueError, errors_impl.InvalidArgumentError), 661 "Shape must be at least rank 1"): 662 op(data=1.0, indices=[0], segment_ids=[3]) 663 664 def testSegmentIdsGreaterThanZero(self): 665 tf_x, np_x = self._input([10, 4], dtype=dtypes_lib.float32) 666 ops_list = [(np.add, None, math_ops.sparse_segment_sum), ( 667 self._mean_cum_op, self._mean_reduce_op, math_ops.sparse_segment_mean)] 668 segment_indices = [1, 2, 2, 2] 669 tf_indices = [8, 3, 0, 9] 670 with self.session(): 671 for np_op1, np_op2, tf_op in ops_list: 672 np_ans = self._sparseSegmentReduce(np_x, tf_indices, segment_indices, 673 np_op1, np_op2) 674 s = tf_op(data=tf_x, indices=tf_indices, segment_ids=segment_indices) 675 tf_ans = self.evaluate(s) 676 self.assertAllClose(np_ans, tf_ans) 677 678 def testValid(self): 679 # Baseline for the test*Invalid* methods below. 680 tf_x, _ = self._input([10, 4], dtype=dtypes_lib.float32) 681 ops_list = [math_ops.sparse_segment_sum, math_ops.sparse_segment_mean] 682 segment_indices = [0, 1, 2, 2] 683 tf_indices = [8, 3, 0, 9] 684 with self.session(): 685 for tf_op in ops_list: 686 s = tf_op(data=tf_x, indices=tf_indices, segment_ids=segment_indices) 687 self.evaluate(s) 688 689 @test_util.run_deprecated_v1 690 def testIndicesInvalid1(self): 691 tf_x, _ = self._input([10, 4], dtype=dtypes_lib.float32) 692 ops_list = [math_ops.sparse_segment_sum, math_ops.sparse_segment_mean] 693 segment_indices = [0, 1, 2, 2] 694 tf_indices = [8, -1, 0, 9] 695 with self.session(use_gpu=False): 696 for tf_op in ops_list: 697 s = tf_op(data=tf_x, indices=tf_indices, segment_ids=segment_indices) 698 with self.assertRaisesOpError( 699 r"indices\[1\] == -1 out of range \[0, 10\)"): 700 self.evaluate(s) 701 702 @test_util.run_deprecated_v1 703 def testIndicesInvalid2(self): 704 tf_x, _ = self._input([10, 4], dtype=dtypes_lib.float32) 705 ops_list = [math_ops.sparse_segment_sum, math_ops.sparse_segment_mean] 706 segment_indices = [0, 1, 2, 2] 707 tf_indices = [8, 3, 0, 10] 708 with self.session(use_gpu=False): 709 for tf_op in ops_list: 710 s = tf_op(data=tf_x, indices=tf_indices, segment_ids=segment_indices) 711 with self.assertRaisesOpError( 712 r"indices\[3\] == 10 out of range \[0, 10\)"): 713 self.evaluate(s) 714 715 @test_util.run_deprecated_v1 716 def testSegmentsInvalid2(self): 717 tf_x, _ = self._input([10, 4], dtype=dtypes_lib.float32) 718 ops_list = [math_ops.sparse_segment_sum, math_ops.sparse_segment_mean] 719 segment_indices = [0, 1, 0, 1] 720 tf_indices = [8, 3, 0, 9] 721 with self.session(use_gpu=False): 722 for tf_op in ops_list: 723 s = tf_op(data=tf_x, indices=tf_indices, segment_ids=segment_indices) 724 with self.assertRaisesOpError("segment ids are not increasing"): 725 self.evaluate(s) 726 727 @test_util.run_deprecated_v1 728 def testSegmentsInvalid3(self): 729 tf_x, _ = self._input([10, 4], dtype=dtypes_lib.float32) 730 ops_list = [math_ops.sparse_segment_sum, math_ops.sparse_segment_mean] 731 segment_indices = [0, 1, 2, 0] 732 tf_indices = [8, 3, 0, 9] 733 with self.session(use_gpu=False): 734 for tf_op in ops_list: 735 s = tf_op(data=tf_x, indices=tf_indices, segment_ids=segment_indices) 736 with self.assertRaisesOpError( 737 r"Segment id 1 out of range \[0, 1\), possibly because " 738 "'segment_ids' input is not sorted"): 739 self.evaluate(s) 740 741 @test_util.run_deprecated_v1 742 def testSegmentsInvalid4(self): 743 tf_x, _ = self._input([10, 4], dtype=dtypes_lib.float32) 744 ops_list = [math_ops.sparse_segment_sum, math_ops.sparse_segment_mean] 745 segment_indices = [-1, 0, 1, 1] 746 tf_indices = [8, 3, 0, 9] 747 with self.session(use_gpu=False): 748 for tf_op in ops_list: 749 s = tf_op(data=tf_x, indices=tf_indices, segment_ids=segment_indices) 750 with self.assertRaisesOpError( 751 r"Segment id -1 out of range \[0, 2\), possibly because " 752 "'segment_ids' input is not sorted"): 753 self.evaluate(s) 754 755 @test_util.run_deprecated_v1 756 def testSegmentsInvalid6(self): 757 tf_x, _ = self._input([10, 4], dtype=dtypes_lib.float32) 758 ops_list = [math_ops.sparse_segment_sum, math_ops.sparse_segment_mean] 759 segment_indices = [0, 0, 0, -1] 760 tf_indices = [8, 3, 0, 9] 761 with self.session(use_gpu=False): 762 for tf_op in ops_list: 763 s = tf_op(data=tf_x, indices=tf_indices, segment_ids=segment_indices) 764 with self.assertRaisesOpError("segment ids must be >= 0"): 765 self.evaluate(s) 766 767 @test_util.run_deprecated_v1 768 def testSegmentsInvalid7(self): 769 tf_x, _ = self._input([10, 4], dtype=dtypes_lib.float32) 770 ops_list = [math_ops.sparse_segment_sum, math_ops.sparse_segment_mean] 771 segment_indices = [0, 0, 0, -2] 772 tf_indices = [8, 3, 0, 9] 773 with self.session(use_gpu=False): 774 for tf_op in ops_list: 775 s = tf_op(data=tf_x, indices=tf_indices, segment_ids=segment_indices) 776 with self.assertRaisesOpError("segment ids must be >= 0"): 777 self.evaluate(s) 778 779 def testSegmentWithNumSegmentsValid(self): 780 # Baseline for the test*WithNumSegmentsInvalid* methods below. 781 tf_x, _ = self._input([10, 4], dtype=dtypes_lib.float32) 782 ops_list = [ 783 math_ops.sparse_segment_sum_with_num_segments, 784 math_ops.sparse_segment_mean_with_num_segments, 785 ] 786 num_segments = 5 787 segment_indices = [0, 1, 3, 3] 788 tf_indices = [8, 3, 0, 9] 789 with self.session(): 790 for tf_op in ops_list: 791 s = tf_op( 792 data=tf_x, 793 indices=tf_indices, 794 segment_ids=segment_indices, 795 num_segments=num_segments) 796 self.evaluate(s) 797 798 @test_util.run_deprecated_v1 799 def testSegmentWithNumSegmentsInvalid1(self): 800 tf_x, _ = self._input([10, 4], dtype=dtypes_lib.float32) 801 ops_list = [ 802 math_ops.sparse_segment_sum_with_num_segments, 803 math_ops.sparse_segment_mean_with_num_segments, 804 ] 805 num_segments = 5 806 segment_indices = [0, 1, 3, 5] 807 tf_indices = [8, 3, 0, 9] 808 with self.session(use_gpu=False): 809 for tf_op in ops_list: 810 s = tf_op( 811 data=tf_x, 812 indices=tf_indices, 813 segment_ids=segment_indices, 814 num_segments=num_segments) 815 with self.assertRaisesOpError("segment ids must be < num_segments"): 816 self.evaluate(s) 817 818 @test_util.run_deprecated_v1 819 def testSegmentWithNumSegmentsInvalid2(self): 820 tf_x, _ = self._input([10, 4], dtype=dtypes_lib.float32) 821 ops_list = [ 822 math_ops.sparse_segment_sum_with_num_segments, 823 math_ops.sparse_segment_mean_with_num_segments, 824 ] 825 num_segments = -2 826 segment_indices = [0, 1, 3, 3] 827 tf_indices = [8, 3, 0, 9] 828 with self.session(use_gpu=False): 829 for tf_op in ops_list: 830 with self.assertRaisesRegex( 831 ValueError, "Cannot specify a negative value for num_segments"): 832 tf_op( 833 data=tf_x, 834 indices=tf_indices, 835 segment_ids=segment_indices, 836 num_segments=num_segments) 837 838 @test_util.run_deprecated_v1 839 def testGradient(self): 840 shape = [10, 4] 841 842 segment_indices = [0, 1, 2, 2] 843 num_indices = len(segment_indices) 844 for tf_op in [math_ops.sparse_segment_sum, math_ops.sparse_segment_mean]: 845 with self.cached_session(): 846 tf_indices, _, tf_x, np_x = self._sparse_input( 847 shape, num_indices, dtype=dtypes_lib.float64) 848 s = tf_op(data=tf_x, indices=tf_indices, segment_ids=segment_indices) 849 jacob_t, jacob_n = gradient_checker.compute_gradient( 850 tf_x, 851 shape, 852 s, [3, 4], 853 x_init_value=np_x.astype(np.double), 854 delta=1) 855 self.assertAllClose(jacob_t, jacob_n) 856 857 @test_util.run_deprecated_v1 858 def testGradientWithEmptySegmentsAtEnd(self): 859 shape = [10, 4] 860 861 num_segments = 5 862 segment_indices = [0, 1, 2, 2] 863 num_indices = len(segment_indices) 864 for tf_op in [ 865 math_ops.sparse_segment_sum_with_num_segments, 866 math_ops.sparse_segment_mean_with_num_segments, 867 ]: 868 with self.cached_session(): 869 tf_indices, _, tf_x, np_x = self._sparse_input( 870 shape, num_indices, dtype=dtypes_lib.float64) 871 s = tf_op( 872 data=tf_x, 873 indices=tf_indices, 874 segment_ids=segment_indices, 875 num_segments=num_segments) 876 jacob_t, jacob_n = gradient_checker.compute_gradient( 877 tf_x, 878 shape, 879 s, [5, 4], 880 x_init_value=np_x.astype(np.double), 881 delta=1) 882 self.assertAllClose(jacob_t, jacob_n) 883 884 def testGradientExplicit(self): 885 # Note that the GPU implem has different paths for different inner sizes. 886 for inner_size in (1, 2, 3, 32): 887 with self.session(): 888 tf_ygrad, np_ygrad = self._input([3, inner_size], 889 dtype=dtypes_lib.float32) 890 segment_ids = [0, 1, 2, 2, 2] 891 indices = [8, 3, 0, 9, 3] 892 output_dim0 = 10 893 ops_list = [ 894 (math_ops.sparse_segment_sum_grad, "sum"), 895 (math_ops.sparse_segment_mean_grad, "mean"), 896 (math_ops.sparse_segment_sqrt_n_grad, "sqrtn"), 897 ] 898 for tf_op, mode in ops_list: 899 np_xgrad = self._sparseSegmentReduceGrad(np_ygrad, indices, 900 segment_ids, output_dim0, 901 mode) 902 tf_xgrad = tf_op(tf_ygrad, indices, segment_ids, output_dim0) 903 self.assertAllClose(tf_xgrad, np_xgrad) 904 905 def testGradientValid(self): 906 # Baseline for the testGradient*Invalid* methods below. 907 tf_x, _ = self._input([3, 4], dtype=dtypes_lib.float32) 908 ops_list = [ 909 math_ops.sparse_segment_sum_grad, math_ops.sparse_segment_mean_grad, 910 math_ops.sparse_segment_sqrt_n_grad 911 ] 912 segment_indices = [0, 1, 2, 2] 913 tf_indices = [8, 3, 0, 9] 914 with self.session(use_gpu=False): 915 for tf_op in ops_list: 916 s = tf_op(tf_x, tf_indices, segment_indices, 10) 917 self.evaluate(s) 918 919 @test_util.run_deprecated_v1 920 def testGradientIndicesInvalid1(self): 921 tf_x, _ = self._input([3, 4], dtype=dtypes_lib.float32) 922 ops_list = [ 923 math_ops.sparse_segment_sum_grad, math_ops.sparse_segment_mean_grad, 924 math_ops.sparse_segment_sqrt_n_grad 925 ] 926 segment_indices = [0, 1, 2, 2] 927 tf_indices = [8, 3, 0, 10] 928 with self.session(use_gpu=False): 929 for tf_op in ops_list: 930 s = tf_op(tf_x, tf_indices, segment_indices, 10) 931 with self.assertRaisesOpError(r"Index 10 out of range \[0, 10\)"): 932 self.evaluate(s) 933 934 @test_util.run_deprecated_v1 935 def testGradientIndicesInvalid2(self): 936 tf_x, _ = self._input([3, 4], dtype=dtypes_lib.float32) 937 ops_list = [ 938 math_ops.sparse_segment_sum_grad, math_ops.sparse_segment_mean_grad, 939 math_ops.sparse_segment_sqrt_n_grad 940 ] 941 segment_indices = [0, 1, 2, 2] 942 tf_indices = [8, 3, -1, 9] 943 with self.session(use_gpu=False): 944 for tf_op in ops_list: 945 s = tf_op(tf_x, tf_indices, segment_indices, 10) 946 with self.assertRaisesOpError(r"Index -1 out of range \[0, 10\)"): 947 self.evaluate(s) 948 949 @test_util.run_deprecated_v1 950 def testGradientSegmentsInvalid1(self): 951 tf_x, _ = self._input( 952 [3, 4], dtype=dtypes_lib.float32) # expecting 3 segments 953 ops_list = [ 954 math_ops.sparse_segment_sum_grad, math_ops.sparse_segment_mean_grad, 955 math_ops.sparse_segment_sqrt_n_grad 956 ] 957 segment_indices = [0, 1, 1, 4] # 5 segments 958 tf_indices = [8, 3, 0, 9] 959 with self.session(use_gpu=False): 960 for tf_op in ops_list: 961 s = tf_op(tf_x, tf_indices, segment_indices, 10) 962 with self.assertRaisesOpError("Invalid number of segments"): 963 self.evaluate(s) 964 965 @test_util.run_deprecated_v1 966 def testGradientSegmentsInvalid2(self): 967 tf_x, _ = self._input([1, 4], dtype=dtypes_lib.float32) 968 ops_list = [ 969 math_ops.sparse_segment_sum_grad, math_ops.sparse_segment_mean_grad, 970 math_ops.sparse_segment_sqrt_n_grad 971 ] 972 segment_indices = [0, 1, 2, 0] 973 tf_indices = [8, 3, 0, 9] 974 with self.session(use_gpu=False): 975 for tf_op in ops_list: 976 s = tf_op(tf_x, tf_indices, segment_indices, 10) 977 with self.assertRaisesOpError(r"Segment id 1 out of range \[0, 1\)"): 978 self.evaluate(s) 979 980 @test_util.run_deprecated_v1 981 def testGradientSegmentsInvalid3(self): 982 tf_x, _ = self._input([2, 4], dtype=dtypes_lib.float32) 983 ops_list = [ 984 math_ops.sparse_segment_sum_grad, math_ops.sparse_segment_mean_grad, 985 math_ops.sparse_segment_sqrt_n_grad 986 ] 987 segment_indices = [-1, 0, 1, 1] 988 tf_indices = [8, 3, 0, 9] 989 with self.session(use_gpu=False): 990 for tf_op in ops_list: 991 s = tf_op(tf_x, tf_indices, segment_indices, 10) 992 with self.assertRaisesOpError(r"Segment id -1 out of range \[0, 2\)"): 993 self.evaluate(s) 994 995 @test_util.run_deprecated_v1 996 def testGradientSegmentsInvalid4(self): 997 tf_x, _ = self._input([0, 4], dtype=dtypes_lib.float32) 998 ops_list = [ 999 math_ops.sparse_segment_sum_grad, math_ops.sparse_segment_mean_grad, 1000 math_ops.sparse_segment_sqrt_n_grad 1001 ] 1002 segment_indices = [0, 1, 2, -1] 1003 tf_indices = [8, 3, 0, 9] 1004 with self.session(use_gpu=False): 1005 for tf_op in ops_list: 1006 s = tf_op(tf_x, tf_indices, segment_indices, 10) 1007 with self.assertRaisesOpError(r"Segment id 0 out of range \[0, 0\)"): 1008 self.evaluate(s) 1009 1010 1011class SegmentReductionOpBenchmark(test.Benchmark): 1012 outer_dim_options = [2**x for x in range(9, 14, 2)] 1013 ratio_options = [2**x for x in range(1, 6, 2)] 1014 inner_dim_options = [2**x for x in range(9, 14, 2)] 1015 # randomly generated sizes with less alignments 1016 inner_dim_options += [ 1017 1120, 1215, 1856, 1302, 1329, 1531, 1313, 1672, 1851, 1584 1018 ] 1019 dtype_options = [np.float32, np.float64] 1020 options = (outer_dim_options, ratio_options, inner_dim_options, dtype_options) 1021 # pylint: disable=g-long-lambda 1022 op_functors = [lambda vc, vs, seg_ids: 1023 ("sorted", math_ops.segment_sum(vc, vs)), 1024 lambda vc, vs, seg_ids: 1025 ("unsorted", 1026 math_ops.unsorted_segment_sum(vc, vs, seg_ids[-1]+1))] 1027 # pylint: enable=g-long-lambda 1028 repeat = 10 1029 1030 def _npTypeToStr(self, t): 1031 if t == np.float32: 1032 return "fp32" 1033 if t == np.float64: 1034 return "fp64" 1035 1036 def _runGraph(self, op_functor, outer_dim, ratio, inner_dim, dtype): 1037 output_outer_dim = int(outer_dim / ratio) 1038 const = np.random.randint(5, size=(outer_dim, inner_dim)) 1039 seg_ids = np.sort(np.random.randint(output_outer_dim, size=outer_dim)) 1040 vs = variables.Variable(seg_ids.astype(np.int32)) 1041 with ops.device("/gpu:0"): 1042 vc = variables.Variable(const.astype(dtype)) 1043 name, op = op_functor(vc, vs, seg_ids) 1044 with session.Session() as sess: 1045 self.evaluate(variables.global_variables_initializer()) 1046 r = self.run_op_benchmark( 1047 sess, 1048 op, 1049 min_iters=self.repeat, 1050 name="_".join( 1051 map(str, 1052 [name, outer_dim, ratio, inner_dim, 1053 self._npTypeToStr(dtype)]))) 1054 return name, r["wall_time"] 1055 1056 def benchmarkSegmentSumGPU(self): 1057 if not test.is_gpu_available(cuda_only=True): 1058 return 1059 for outer_dim, ratio, inner_dim, dtype in itertools.product(*self.options): 1060 op_functor = self.op_functors[0] 1061 with ops.Graph().as_default(): 1062 self._runGraph(op_functor, outer_dim, ratio, inner_dim, dtype) 1063 1064 def benchmarkUnsortedSegmentSumGPU(self): 1065 if not test.is_gpu_available(cuda_only=True): 1066 return 1067 for outer_dim, ratio, inner_dim, dtype in itertools.product(*self.options): 1068 op_functor = self.op_functors[1] 1069 with ops.Graph().as_default(): 1070 self._runGraph(op_functor, outer_dim, ratio, inner_dim, dtype) 1071 1072 1073if __name__ == "__main__": 1074 test.main() 1075