| # 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. |
| # ============================================================================== |
| |
| from absl.testing import parameterized |
| import numpy as np |
| |
| from tensorflow.compiler.tests import xla_test |
| from tensorflow.python.framework import constant_op |
| from tensorflow.python.framework import dtypes |
| from tensorflow.python.ops import array_ops |
| from tensorflow.python.platform import test |
| |
| |
| class MatrixBandPartTest(xla_test.XLATestCase, parameterized.TestCase): |
| |
| @parameterized.parameters( |
| { |
| 'batch_shape': [], |
| 'rows': 1, |
| 'cols': 1 |
| }, |
| { |
| 'batch_shape': [], |
| 'rows': 1, |
| 'cols': 2 |
| }, |
| { |
| 'batch_shape': [], |
| 'rows': 1, |
| 'cols': 7 |
| }, |
| { |
| 'batch_shape': [], |
| 'rows': 2, |
| 'cols': 1 |
| }, |
| { |
| 'batch_shape': [], |
| 'rows': 2, |
| 'cols': 2 |
| }, |
| { |
| 'batch_shape': [], |
| 'rows': 2, |
| 'cols': 7 |
| }, |
| { |
| 'batch_shape': [], |
| 'rows': 7, |
| 'cols': 1 |
| }, |
| { |
| 'batch_shape': [], |
| 'rows': 7, |
| 'cols': 2 |
| }, |
| { |
| 'batch_shape': [], |
| 'rows': 7, |
| 'cols': 7 |
| }, |
| { |
| 'batch_shape': [2,], |
| 'rows': 1, |
| 'cols': 1 |
| }, |
| { |
| 'batch_shape': [2,], |
| 'rows': 1, |
| 'cols': 2 |
| }, |
| { |
| 'batch_shape': [2,], |
| 'rows': 1, |
| 'cols': 7 |
| }, |
| { |
| 'batch_shape': [2,], |
| 'rows': 2, |
| 'cols': 1 |
| }, |
| { |
| 'batch_shape': [2,], |
| 'rows': 2, |
| 'cols': 2 |
| }, |
| { |
| 'batch_shape': [2,], |
| 'rows': 2, |
| 'cols': 7 |
| }, |
| { |
| 'batch_shape': [2,], |
| 'rows': 7, |
| 'cols': 1 |
| }, |
| { |
| 'batch_shape': [2,], |
| 'rows': 7, |
| 'cols': 2 |
| }, |
| { |
| 'batch_shape': [2,], |
| 'rows': 7, |
| 'cols': 7 |
| }, |
| { |
| 'batch_shape': [1, 3, 2], |
| 'rows': 1, |
| 'cols': 1 |
| }, |
| { |
| 'batch_shape': [1, 3, 2], |
| 'rows': 1, |
| 'cols': 2 |
| }, |
| { |
| 'batch_shape': [1, 3, 2], |
| 'rows': 1, |
| 'cols': 7 |
| }, |
| { |
| 'batch_shape': [1, 3, 2], |
| 'rows': 2, |
| 'cols': 1 |
| }, |
| { |
| 'batch_shape': [1, 3, 2], |
| 'rows': 2, |
| 'cols': 2 |
| }, |
| { |
| 'batch_shape': [1, 3, 2], |
| 'rows': 2, |
| 'cols': 7 |
| }, |
| { |
| 'batch_shape': [1, 3, 2], |
| 'rows': 7, |
| 'cols': 1 |
| }, |
| { |
| 'batch_shape': [1, 3, 2], |
| 'rows': 7, |
| 'cols': 2 |
| }, |
| { |
| 'batch_shape': [1, 3, 2], |
| 'rows': 7, |
| 'cols': 7 |
| }, |
| ) |
| def testMatrixBandPart(self, batch_shape, rows, cols): |
| # TODO(b/125505881): Disabled due to LLVM backend crash. |
| if self.device == 'XLA_CPU' and cols == 7 and rows == 1 and batch_shape == [ |
| 1, 3, 2 |
| ]: |
| pass |
| for dtype in self.float_types: |
| with self.session(): |
| mat = np.ones(batch_shape + [rows, cols]).astype(dtype) |
| batch_mat = np.tile(mat, batch_shape + [1, 1]) |
| for lower in -1, 0, 1, rows - 1: |
| for upper in -1, 0, 1, cols - 1: |
| band_np = mat |
| if lower >= 0: |
| band_np = np.triu(band_np, -lower) |
| if upper >= 0: |
| band_np = np.tril(band_np, upper) |
| if batch_shape: |
| band_np = np.tile(band_np, batch_shape + [1, 1]) |
| |
| placeholder = array_ops.placeholder(dtype) |
| with self.test_scope(): |
| band = array_ops.matrix_band_part( |
| placeholder, constant_op.constant(lower, dtype=dtypes.int32), |
| constant_op.constant(upper, dtype=dtypes.int32)) |
| feed_dict = {placeholder: batch_mat} |
| self.assertAllEqual(band_np, band.eval(feed_dict=feed_dict)) |
| |
| |
| if __name__ == "__main__": |
| test.main() |