# Copyright 2019 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 XLA matrix diag ops.""" import numpy as np from tensorflow.compiler.tests import xla_test from tensorflow.python.ops import array_ops from tensorflow.python.ops import gen_array_ops from tensorflow.python.platform import googletest default_v2_alignment = "LEFT_LEFT" alignment_list = ["RIGHT_LEFT", "LEFT_RIGHT"] def zip_to_first_list_length(a, b): if len(b) > len(a): return zip(a, b[:len(a)]) return zip(a, b + [None] * (len(a) - len(b))) # Routines to convert test cases to have diagonals in a specified alignment. # Copied from //third_party/tensorflow/python/kernel_tests/array_ops/ # diag_op_test.py def repack_diagonals(packed_diagonals, diag_index, num_rows, num_cols, align=None): # The original test cases are LEFT_LEFT aligned. if align == default_v2_alignment or align is None: return packed_diagonals align = align.split("_") d_lower, d_upper = diag_index batch_dims = packed_diagonals.ndim - (2 if d_lower < d_upper else 1) max_diag_len = packed_diagonals.shape[-1] index = (slice(None),) * batch_dims repacked_diagonals = np.zeros_like(packed_diagonals) # Aligns each diagonal row-by-row. for diag_index in range(d_lower, d_upper + 1): diag_len = min(num_rows + min(0, diag_index), num_cols - max(0, diag_index)) row_index = d_upper - diag_index padding_len = max_diag_len - diag_len left_align = (diag_index >= 0 and align[0] == "LEFT") or (diag_index <= 0 and align[1] == "LEFT") # Prepares index tuples. extra_dim = tuple() if d_lower == d_upper else (row_index,) packed_last_dim = (slice(None),) if left_align else (slice(0, diag_len, 1),) repacked_last_dim = (slice(None),) if left_align else (slice( padding_len, max_diag_len, 1),) packed_index = index + extra_dim + packed_last_dim repacked_index = index + extra_dim + repacked_last_dim # Repacks the diagonal. repacked_diagonals[repacked_index] = packed_diagonals[packed_index] return repacked_diagonals def repack_diagonals_in_tests(tests, align=None): # The original test cases are LEFT_LEFT aligned. if align == default_v2_alignment or align is None: return tests new_tests = dict() # Loops through each case. for diag_index, (packed_diagonals, padded_diagonals) in tests.items(): num_rows, num_cols = padded_diagonals.shape[-2:] repacked_diagonals = repack_diagonals( packed_diagonals, diag_index, num_rows, num_cols, align=align) new_tests[diag_index] = (repacked_diagonals, padded_diagonals) return new_tests # Test cases shared by MatrixDiagV2, MatrixDiagPartV2, and MatrixSetDiagV2. # Copied from //third_party/tensorflow/python/kernel_tests/array_ops/ # diag_op_test.py def square_cases(align=None): # pyformat: disable mat = np.array([[[1, 2, 3, 4, 5], [6, 7, 8, 9, 1], [3, 4, 5, 6, 7], [8, 9, 1, 2, 3], [4, 5, 6, 7, 8]], [[9, 1, 2, 3, 4], [5, 6, 7, 8, 9], [1, 2, 3, 4, 5], [6, 7, 8, 9, 1], [2, 3, 4, 5, 6]]]) tests = dict() # tests[d_lower, d_upper] = (compact_diagonals, padded_diagonals) tests[-1, -1] = (np.array([[6, 4, 1, 7], [5, 2, 8, 5]]), np.array([[[0, 0, 0, 0, 0], [6, 0, 0, 0, 0], [0, 4, 0, 0, 0], [0, 0, 1, 0, 0], [0, 0, 0, 7, 0]], [[0, 0, 0, 0, 0], [5, 0, 0, 0, 0], [0, 2, 0, 0, 0], [0, 0, 8, 0, 0], [0, 0, 0, 5, 0]]])) tests[-4, -3] = (np.array([[[8, 5], [4, 0]], [[6, 3], [2, 0]]]), np.array([[[0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [8, 0, 0, 0, 0], [4, 5, 0, 0, 0]], [[0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [6, 0, 0, 0, 0], [2, 3, 0, 0, 0]]])) tests[-2, 1] = (np.array([[[2, 8, 6, 3, 0], [1, 7, 5, 2, 8], [6, 4, 1, 7, 0], [3, 9, 6, 0, 0]], [[1, 7, 4, 1, 0], [9, 6, 3, 9, 6], [5, 2, 8, 5, 0], [1, 7, 4, 0, 0]]]), np.array([[[1, 2, 0, 0, 0], [6, 7, 8, 0, 0], [3, 4, 5, 6, 0], [0, 9, 1, 2, 3], [0, 0, 6, 7, 8]], [[9, 1, 0, 0, 0], [5, 6, 7, 0, 0], [1, 2, 3, 4, 0], [0, 7, 8, 9, 1], [0, 0, 4, 5, 6]]])) tests[2, 4] = (np.array([[[5, 0, 0], [4, 1, 0], [3, 9, 7]], [[4, 0, 0], [3, 9, 0], [2, 8, 5]]]), np.array([[[0, 0, 3, 4, 5], [0, 0, 0, 9, 1], [0, 0, 0, 0, 7], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0]], [[0, 0, 2, 3, 4], [0, 0, 0, 8, 9], [0, 0, 0, 0, 5], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0]]])) # pyformat: enable return (mat, repack_diagonals_in_tests(tests, align)) def tall_cases(align=None): # pyformat: disable mat = np.array([[[1, 2, 3], [4, 5, 6], [7, 8, 9], [9, 8, 7], [6, 5, 4]], [[3, 2, 1], [1, 2, 3], [4, 5, 6], [7, 8, 9], [9, 8, 7]]]) tests = dict() # tests[d_lower, d_upper] = (compact_diagonals, padded_diagonals) tests[0, 0] = (np.array([[1, 5, 9], [3, 2, 6]]), np.array([[[1, 0, 0], [0, 5, 0], [0, 0, 9], [0, 0, 0]], [[3, 0, 0], [0, 2, 0], [0, 0, 6], [0, 0, 0]]])) tests[-4, -3] = (np.array([[[9, 5], [6, 0]], [[7, 8], [9, 0]]]), np.array([[[0, 0, 0], [0, 0, 0], [0, 0, 0], [9, 0, 0], [6, 5, 0]], [[0, 0, 0], [0, 0, 0], [0, 0, 0], [7, 0, 0], [9, 8, 0]]])) tests[-2, -1] = (np.array([[[4, 8, 7], [7, 8, 4]], [[1, 5, 9], [4, 8, 7]]]), np.array([[[0, 0, 0], [4, 0, 0], [7, 8, 0], [0, 8, 7], [0, 0, 4]], [[0, 0, 0], [1, 0, 0], [4, 5, 0], [0, 8, 9], [0, 0, 7]]])) tests[-2, 1] = (np.array([[[2, 6, 0], [1, 5, 9], [4, 8, 7], [7, 8, 4]], [[2, 3, 0], [3, 2, 6], [1, 5, 9], [4, 8, 7]]]), np.array([[[1, 2, 0], [4, 5, 6], [7, 8, 9], [0, 8, 7], [0, 0, 4]], [[3, 2, 0], [1, 2, 3], [4, 5, 6], [0, 8, 9], [0, 0, 7]]])) tests[1, 2] = (np.array([[[3, 0], [2, 6]], [[1, 0], [2, 3]]]), np.array([[[0, 2, 3], [0, 0, 6], [0, 0, 0], [0, 0, 0], [0, 0, 0]], [[0, 2, 1], [0, 0, 3], [0, 0, 0], [0, 0, 0], [0, 0, 0]]])) # pyformat: enable return (mat, repack_diagonals_in_tests(tests, align)) def fat_cases(align=None): # pyformat: disable mat = np.array([[[1, 2, 3, 4], [5, 6, 7, 8], [9, 1, 2, 3]], [[4, 5, 6, 7], [8, 9, 1, 2], [3, 4, 5, 6]]]) tests = dict() # tests[d_lower, d_upper] = (compact_diagonals, padded_diagonals) tests[0, 0] = (np.array([[1, 6, 2], [4, 9, 5]]), np.array([[[1, 0, 0, 0], [0, 6, 0, 0], [0, 0, 2, 0]], [[4, 0, 0, 0], [0, 9, 0, 0], [0, 0, 5, 0]]])) tests[2, 2] = (np.array([[3, 8], [6, 2]]), np.array([[[0, 0, 3, 0], [0, 0, 0, 8], [0, 0, 0, 0]], [[0, 0, 6, 0], [0, 0, 0, 2], [0, 0, 0, 0]]])) tests[-2, 0] = (np.array([[[1, 6, 2], [5, 1, 0], [9, 0, 0]], [[4, 9, 5], [8, 4, 0], [3, 0, 0]]]), np.array([[[1, 0, 0, 0], [5, 6, 0, 0], [9, 1, 2, 0]], [[4, 0, 0, 0], [8, 9, 0, 0], [3, 4, 5, 0]]])) tests[-1, 1] = (np.array([[[2, 7, 3], [1, 6, 2], [5, 1, 0]], [[5, 1, 6], [4, 9, 5], [8, 4, 0]]]), np.array([[[1, 2, 0, 0], [5, 6, 7, 0], [0, 1, 2, 3]], [[4, 5, 0, 0], [8, 9, 1, 0], [0, 4, 5, 6]]])) tests[0, 3] = (np.array([[[4, 0, 0], [3, 8, 0], [2, 7, 3], [1, 6, 2]], [[7, 0, 0], [6, 2, 0], [5, 1, 6], [4, 9, 5]]]), np.array([[[1, 2, 3, 4], [0, 6, 7, 8], [0, 0, 2, 3]], [[4, 5, 6, 7], [0, 9, 1, 2], [0, 0, 5, 6]]])) # pyformat: enable return (mat, repack_diagonals_in_tests(tests, align)) def all_tests(align=None): return [square_cases(align), tall_cases(align), fat_cases(align)] class MatrixDiagTest(xla_test.XLATestCase): def _assertOpOutputMatchesExpected(self, params, solution, high_level=True, rtol=1e-3, atol=1e-5): """Verifies that matrix_diag produces `solution` when fed `params`. Args: params: dictionary containing input parameters to matrix_diag. solution: numpy array representing the expected output of matrix_diag. high_level: call high_level matrix_diag rtol: relative tolerance for equality test. atol: absolute tolerance for equality test. """ diagonal = params["diagonal"] with self.session() as session: for dtype in self.numeric_types - {np.int8, np.uint8}: expected = solution.astype(dtype) with self.test_scope(): params["diagonal"] = array_ops.placeholder( dtype, diagonal.shape, name="diagonal") if high_level: # wraps gen_array_ops.matrix_diag_v3 output = array_ops.matrix_diag(**params) else: # TODO(b/201086188): Remove this case once MatrixDiag V1 is removed. output = gen_array_ops.matrix_diag(**params) result = session.run(output, {params["diagonal"]: diagonal.astype(dtype)}) self.assertEqual(output.dtype, expected.dtype) self.assertAllCloseAccordingToType( expected, result, rtol=rtol, atol=atol, bfloat16_rtol=0.03) # Generic tests applicable to both v1 and v2 ops. # Originally from unary_ops_tests.py. def _testV1Level(self, high_level): # pyformat: disable vecs1 = np.array([[1, 2], [3, 4]]) solution1 = np.array([[[1, 0], [0, 2]], [[3, 0], [0, 4]]]) vecs2 = np.array([1, 2, 3, 4]) solution2 = np.array([[1, 0, 0, 0], [0, 2, 0, 0], [0, 0, 3, 0], [0, 0, 0, 4]]) vecs3 = np.array([[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], # pylint: disable=bad-whitespace [10, 11, 12]]]) solution3 = np.array([[[[1, 0, 0], [0, 2, 0], [0, 0, 3]], [[4, 0, 0], [0, 5, 0], [0, 0, 6]]], [[[7, 0, 0], [0, 8, 0], [0, 0, 9]], [[10, 0, 0], [0, 11, 0], [0, 0, 12]]]]) # pyformat: enable self._assertOpOutputMatchesExpected({"diagonal": vecs1}, solution1, high_level) self._assertOpOutputMatchesExpected({"diagonal": vecs2}, solution2, high_level) self._assertOpOutputMatchesExpected({"diagonal": vecs3}, solution3, high_level) def testV1(self): self._testV1Level(True) def testV1LowLevel(self): self._testV1Level(False) # From here onwards are v2-only tests. def testSquare(self): for align in alignment_list: for _, tests in [square_cases(align)]: for diag_index, (vecs, solution) in tests.items(): params = {"diagonal": vecs[0], "k": diag_index, "align": align} self._assertOpOutputMatchesExpected(params, solution[0]) def testSquareBatch(self): for align in alignment_list: for _, tests in [square_cases(align)]: for diag_index, (vecs, solution) in tests.items(): params = {"diagonal": vecs, "k": diag_index, "align": align} self._assertOpOutputMatchesExpected(params, solution) def testRectangularBatch(self): # Stores expected num_rows and num_cols (when the other is given). # expected[(d_lower, d_upper)] = (expected_num_rows, expected_num_cols) test_list = list() # Do not align the test cases here. Re-alignment needs to happen after the # solution shape is updated. # Square cases: expected = { (-1, -1): (5, 4), (-4, -3): (5, 2), (-2, 1): (5, 5), (2, 4): (3, 5), } test_list.append((expected, square_cases())) # Tall cases expected = { (0, 0): (3, 3), (-4, -3): (5, 2), (-2, -1): (4, 3), (-2, 1): (3, 3), (1, 2): (2, 3) } test_list.append((expected, tall_cases())) # Fat cases expected = { (2, 2): (2, 4), (-2, 0): (3, 3), (-1, 1): (3, 3), (0, 3): (3, 3) } test_list.append((expected, fat_cases())) # Giving both num_rows and num_cols align = alignment_list[0] for _, tests in [tall_cases(align), fat_cases(align)]: for diag_index, (vecs, solution) in tests.items(): self._assertOpOutputMatchesExpected( { "diagonal": vecs, "k": diag_index, "num_rows": solution.shape[-2], "num_cols": solution.shape[-1], "align": align }, solution) # We go through each alignment in a round-robin manner. align_index = 0 # Giving just num_rows or num_cols. for expected, (_, tests) in test_list: for diag_index, (new_num_rows, new_num_cols) in expected.items(): align = alignment_list[align_index] align_index = (align_index + 1) % len(alignment_list) vecs, solution = tests[diag_index] solution_given_num_rows = solution.take( indices=range(new_num_cols), axis=-1) # Repacks the diagonal input according to the new solution shape. vecs_given_num_rows = repack_diagonals( vecs, diag_index, solution_given_num_rows.shape[-2], new_num_cols, align=align) self._assertOpOutputMatchesExpected( { "diagonal": vecs_given_num_rows, "k": diag_index, "num_rows": solution_given_num_rows.shape[-2], "align": align }, solution_given_num_rows) solution_given_num_cols = solution.take( indices=range(new_num_rows), axis=-2) # Repacks the diagonal input according to the new solution shape. vecs_given_num_cols = repack_diagonals( vecs, diag_index, new_num_rows, solution_given_num_cols.shape[-1], align=align) self._assertOpOutputMatchesExpected( { "diagonal": vecs_given_num_cols, "k": diag_index, "num_cols": solution_given_num_cols.shape[-1], "align": align }, solution_given_num_cols) def testPadding(self): for padding_value, align in zip_to_first_list_length([555, -11], alignment_list): for _, tests in all_tests(align): for diag_index, (vecs, solution) in tests.items(): mask = (solution == 0) solution = solution + (mask * padding_value) self._assertOpOutputMatchesExpected( { "diagonal": vecs, "k": diag_index, "num_rows": solution.shape[-2], "num_cols": solution.shape[-1], "padding_value": padding_value, "align": align }, solution) class MatrixSetDiagTest(xla_test.XLATestCase): def _assertOpOutputMatchesExpected(self, params, solution, high_level=True, rtol=1e-3, atol=1e-5): """Verifies that matrix_set_diag produces `solution` when fed `params`. Args: params: dictionary containing input parameters to matrix_set_diag. solution: numpy array representing the expected output of matrix_set_diag. high_level: call high_level matrix_set_diag rtol: relative tolerance for equality test. atol: absolute tolerance for equality test. """ input = params["input"] # pylint: disable=redefined-builtin diagonal = params["diagonal"] with self.session() as session: for dtype in self.numeric_types - {np.int8, np.uint8}: expected = solution.astype(dtype) with self.test_scope(): params["input"] = array_ops.placeholder( dtype, input.shape, name="input") params["diagonal"] = array_ops.placeholder( dtype, diagonal.shape, name="diagonal") if high_level: # wraps gen_array_ops.matrix_set_diag_v3 output = array_ops.matrix_set_diag(**params) else: # TODO(b/201086188): Remove this case once MatrixDiag V1 is removed. output = gen_array_ops.matrix_set_diag(**params) result = session.run( output, { params["input"]: input.astype(dtype), params["diagonal"]: diagonal.astype(dtype) }) self.assertEqual(output.dtype, expected.dtype) self.assertAllCloseAccordingToType( expected, result, rtol=rtol, atol=atol, bfloat16_rtol=0.03) # Generic tests applicable to both v1 and v2 ops. # Originally from binary_ops_tests.py. def _testV1Level(self, high_level): test_cases = list() # pyformat: disable # pylint: disable=bad-whitespace # Square cases. input = np.array([[0, 1, 0], # pylint: disable=redefined-builtin [1, 0, 1], [1, 1, 1]]) diag = np.array([1, 2, 3]) solution = np.array([[1, 1, 0], [1, 2, 1], [1, 1, 3]]) test_cases.append(({"input": input, "diagonal": diag}, solution)) input = np.array([[[1, 0, 3], [0, 2, 0], [1, 0, 3]], [[4, 0, 4], [0, 5, 0], [2, 0, 6]]]) diag = np.array([[-1, 0, -3], [-4, -5, -6]]) solution = np.array([[[-1, 0, 3], [ 0, 0, 0], [ 1, 0, -3]], [[-4, 0, 4], [ 0, -5, 0], [ 2, 0, -6]]]) test_cases.append(({"input": input, "diagonal": diag}, solution)) # Rectangular cases. input = np.array([[0, 1, 0], [1, 0, 1]]) diag = np.array([3, 4]) solution = np.array([[3, 1, 0], [1, 4, 1]]) test_cases.append(({"input": input, "diagonal": diag}, solution)) input = np.array([[0, 1], [1, 0], [1, 1]]) diag = np.array([3, 4]) solution = np.array([[3, 1], [1, 4], [1, 1]]) test_cases.append(({"input": input, "diagonal": diag}, solution)) input = np.array([[[1, 0, 3], [0, 2, 0]], [[4, 0, 4], [0, 5, 0]]]) diag = np.array([[-1, -2], [-4, -5]]) solution = np.array([[[-1, 0, 3], [ 0, -2, 0]], [[-4, 0, 4], [ 0, -5, 0]]]) test_cases.append(({"input": input, "diagonal": diag}, solution)) # pylint: enable=bad-whitespace # pyformat: enable for test in test_cases: self._assertOpOutputMatchesExpected(test[0], test[1], high_level) def testV1(self): self._testV1Level(True) def testV1LowLevel(self): self._testV1Level(False) # From here onwards are v2-only tests. def testSingleMatrix(self): for align in alignment_list: for _, tests in all_tests(align): for diag_index, (vecs, banded_mat) in tests.items(): mask = (banded_mat[0] == 0) input_mat = np.random.randint(10, size=mask.shape) solution = input_mat * mask + banded_mat[0] self._assertOpOutputMatchesExpected( { "input": input_mat, "diagonal": vecs[0], "k": diag_index, "align": align }, solution) def testBatch(self): for align in alignment_list: for _, tests in all_tests(align): for diag_index, (vecs, banded_mat) in tests.items(): mask = (banded_mat == 0) input_mat = np.random.randint(10, size=mask.shape) solution = input_mat * mask + banded_mat self._assertOpOutputMatchesExpected( { "input": input_mat, "diagonal": vecs, "k": diag_index, "align": align }, solution) class MatrixDiagPartTest(xla_test.XLATestCase): def _assertOpOutputMatchesExpected(self, params, solution, high_level=True, rtol=1e-3, atol=1e-5): """Verifies that matrix_diag_part produces `solution` when fed `params`. Args: params: dictionary containing input parameters to matrix_diag_part. solution: numpy array representing the expected output. high_level: call high_level matrix_set_diag rtol: relative tolerance for equality test. atol: absolute tolerance for equality test. """ input = params["input"] # pylint: disable=redefined-builtin with self.session() as session: for dtype in self.numeric_types - {np.int8, np.uint8}: expected = solution.astype(dtype) with self.test_scope(): params["input"] = array_ops.placeholder( dtype, input.shape, name="input") if high_level: # wraps gen_array_ops.matrix_diag_part_v3 output = array_ops.matrix_diag_part(**params) else: # TODO(b/201086188): Remove this case once MatrixDiag V1 is removed. output = gen_array_ops.matrix_diag_part(**params) output = array_ops.matrix_diag_part(**params) result = session.run(output, { params["input"]: input.astype(dtype), }) self.assertEqual(output.dtype, expected.dtype) self.assertAllCloseAccordingToType( expected, result, rtol=rtol, atol=atol, bfloat16_rtol=0.03) # Generic tests applicable to both v1 and v2 ops. # Originally from unary_ops_tests.py. def _testV1Level(self, high_level): matrices = np.arange(3 * 2 * 4).reshape([3, 2, 4]) solution = np.array([[0, 5], [8, 13], [16, 21]]) self._assertOpOutputMatchesExpected({"input": matrices}, solution, high_level) def testV1(self): self._testV1Level(True) def testV1LowLevel(self): self._testV1Level(False) # From here onwards are v2-only tests. def testSingleMatrix(self): for align in alignment_list: test_list = [square_cases(align), tall_cases(align), fat_cases(align)] for mat, tests in test_list: for diag_index, (solution, _) in tests.items(): self._assertOpOutputMatchesExpected( { "input": mat[0], "k": diag_index, "align": align }, solution[0]) def testBatch(self): for align in alignment_list: for mat, tests in all_tests(align): for diag_index, (solution, _) in tests.items(): self._assertOpOutputMatchesExpected( { "input": mat, "k": diag_index, "align": align }, solution) def testPadding(self): for padding_value, align in zip_to_first_list_length([555, -11], alignment_list): for mat, tests in all_tests(align): for diag_index, (solution, _) in tests.items(): mask = (solution == 0) solution = solution + (mask * padding_value) self._assertOpOutputMatchesExpected( { "input": mat, "k": diag_index, "padding_value": padding_value, "align": align }, solution) if __name__ == "__main__": googletest.main()