| # Copyright 2022 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 tridiagonal_matmul_ops.""" |
| |
| import numpy as np |
| |
| from tensorflow.compiler.tests import xla_test |
| from tensorflow.python.eager import def_function |
| from tensorflow.python.framework import constant_op |
| from tensorflow.python.framework import dtypes |
| from tensorflow.python.framework import ops |
| from tensorflow.python.ops import array_ops |
| from tensorflow.python.ops import array_ops_stack # pylint: disable=g-direct-tensorflow-import |
| from tensorflow.python.ops import gradient_checker_v2 |
| from tensorflow.python.ops import math_ops |
| from tensorflow.python.ops import stateless_random_ops as random |
| from tensorflow.python.ops.linalg import linalg_impl |
| from tensorflow.python.platform import googletest |
| |
| |
| class TridiagonalMatMulOpsTest(xla_test.XLATestCase): |
| |
| @def_function.function(jit_compile=True) |
| def _jit_tridiagonal_matmul(self, diagonals, rhs, diagonals_format): |
| return linalg_impl.tridiagonal_matmul( |
| diagonals, rhs, diagonals_format=diagonals_format) |
| |
| @def_function.function(jit_compile=False) |
| def _tridiagonal_matmul(self, diagonals, rhs, diagonals_format): |
| return linalg_impl.tridiagonal_matmul( |
| diagonals, rhs, diagonals_format=diagonals_format) |
| |
| def _testAllFormats(self, |
| superdiag, |
| maindiag, |
| subdiag, |
| rhs, |
| dtype=dtypes.float64): |
| superdiag_extended = np.pad(superdiag, [0, 1], 'constant') |
| subdiag_extended = np.pad(subdiag, [1, 0], 'constant') |
| |
| diags_compact = np.stack([superdiag_extended, maindiag, subdiag_extended]) |
| diags_matrix = np.diag(superdiag, 1) + np.diag(maindiag, 0) + np.diag( |
| subdiag, -1) |
| |
| with self.test_scope(): |
| diags_sequence_op = (constant_op.constant(superdiag_extended, dtype), |
| constant_op.constant(maindiag, dtype), |
| constant_op.constant(subdiag_extended, dtype)) |
| diags_compact_op = constant_op.constant(diags_compact, dtype) |
| diags_matrix_op = constant_op.constant(diags_matrix, dtype) |
| rhs_op = constant_op.constant(rhs, dtype) |
| |
| rhs_batch_op = array_ops_stack.stack([rhs_op, 2 * rhs_op]) |
| diags_compact_batch_op = array_ops_stack.stack( |
| [diags_compact_op, 2 * diags_compact_op]) |
| diags_matrix_batch_op = array_ops_stack.stack( |
| [diags_matrix_op, 2 * diags_matrix_op]) |
| diags_sequence_batch_op = [ |
| array_ops_stack.stack([x, 2 * x]) for x in diags_sequence_op |
| ] |
| |
| results = [ |
| self._jit_tridiagonal_matmul( |
| diags_sequence_op, rhs_op, diagonals_format='sequence'), |
| self._jit_tridiagonal_matmul( |
| diags_compact_op, rhs_op, diagonals_format='compact'), |
| self._jit_tridiagonal_matmul( |
| diags_matrix_op, rhs_op, diagonals_format='matrix') |
| ] |
| |
| results_batch = [ |
| self._jit_tridiagonal_matmul( |
| diags_sequence_batch_op, |
| rhs_batch_op, |
| diagonals_format='sequence'), |
| self._jit_tridiagonal_matmul( |
| diags_compact_batch_op, rhs_batch_op, diagonals_format='compact'), |
| self._jit_tridiagonal_matmul( |
| diags_matrix_batch_op, rhs_batch_op, diagonals_format='matrix') |
| ] |
| |
| expected = self._tridiagonal_matmul( |
| diags_sequence_op, rhs_op, diagonals_format='sequence') |
| expected_batch = np.stack([expected, 4 * expected]) |
| for result in results: |
| self.assertAllClose(result, expected) |
| for result in results_batch: |
| self.assertAllClose(result, expected_batch) |
| |
| def _makeTridiagonalMatrix(self, superdiag, maindiag, subdiag): |
| super_pad = [[0, 0], [0, 1], [1, 0]] |
| sub_pad = [[0, 0], [1, 0], [0, 1]] |
| |
| super_part = array_ops.pad(array_ops.matrix_diag(superdiag), super_pad) |
| main_part = array_ops.matrix_diag(maindiag) |
| sub_part = array_ops.pad(array_ops.matrix_diag(subdiag), sub_pad) |
| return super_part + main_part + sub_part |
| |
| def _gradientTest(self, diags, rhs, dtype=dtypes.float64): |
| |
| def reference_matmul(diags, rhs): |
| matrix = self._makeTridiagonalMatrix(diags[..., 0, :-1], diags[..., 1, :], |
| diags[..., 2, 1:]) |
| return math_ops.matmul(matrix, rhs) |
| |
| with self.session(), self.test_scope(): |
| diags = constant_op.constant(diags, dtype=dtype) |
| rhs = constant_op.constant(rhs, dtype=dtype) |
| grad_reference, _ = gradient_checker_v2.compute_gradient( |
| reference_matmul, [diags, rhs]) |
| grad_theoretical, grad_numerical = gradient_checker_v2.compute_gradient( |
| linalg_impl.tridiagonal_matmul, [diags, rhs]) |
| self.assertAllClose(grad_theoretical, grad_numerical) |
| self.assertAllClose(grad_theoretical, grad_reference) |
| |
| def test2x2(self): |
| self._testAllFormats( |
| superdiag=[1], maindiag=[2, 3], subdiag=[4], rhs=[[2, 1], [4, 3]]) |
| |
| def test3x3(self): |
| for dtype in [dtypes.float32, dtypes.float64]: |
| self._testAllFormats( |
| superdiag=[1, 2], |
| maindiag=[1, 2, 1], |
| subdiag=[2, 1], |
| rhs=[[1, 1], [2, 2], [3, 3]], |
| dtype=dtype) |
| |
| def testNDimensional(self): |
| with self.test_scope(): |
| shape = [77, 10, 3, 5, 7] |
| superdiag = random.stateless_random_uniform( |
| shape=shape[:-1], dtype=dtypes.float32, seed=[5, 10]) |
| maindiag = random.stateless_random_uniform( |
| shape=shape[:-1], dtype=dtypes.float32, seed=[5, 11]) |
| subdiag = random.stateless_random_uniform( |
| shape=shape[:-1], dtype=dtypes.float32, seed=[5, 12]) |
| rhs = random.stateless_random_uniform( |
| shape=shape, dtype=dtypes.float32, seed=[5, 13]) |
| |
| expected = self._tridiagonal_matmul((superdiag, maindiag, subdiag), |
| rhs, |
| diagonals_format='sequence') |
| |
| real = self._jit_tridiagonal_matmul((superdiag, maindiag, subdiag), |
| rhs, |
| diagonals_format='sequence') |
| |
| self.assertAllClose(expected, real) |
| |
| def testGradientSmall(self): |
| self._gradientTest([[[1, 2, 0], [1, 2, 3], [0, 1, 2]]], |
| [[[1, 2], [3, 4], [5, 6]]], |
| dtype=dtypes.float64) |
| |
| # testComplex is skipped as complex type is not yet supported. |
| |
| |
| if __name__ == '__main__': |
| ops.enable_eager_execution() |
| googletest.main() |
