From 1130d49032fba29a3308bef6f82f27ecf6c422e0 Mon Sep 17 00:00:00 2001 From: Joe Evans Date: Fri, 24 Jul 2020 09:50:46 -0700 Subject: [PATCH] Fix linalg_potri and linalg_potrf operators for large tensor. (#18752) * Fix linalg_potri operator for large tensor. * Update other variables to support large tensors. * Add to contributors. * Fix whitespace. * Update ZeroTriangular to support large tensors. * Add large tensor unit tests for linalg_potrf and linalg_potri. * Fix crash when accessing already destructed static variables (#18768) (#18778) Co-authored-by: Joe Evans Co-authored-by: Przemyslaw Tredak --- CONTRIBUTORS.md | 1 + src/operator/tensor/la_op-inl.h | 11 ++++++----- tests/nightly/test_large_array.py | 27 +++++++++++++++++++++++++++ 3 files changed, 34 insertions(+), 5 deletions(-) diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md index 81b20d4b092d..bd7f966aaa5f 100644 --- a/CONTRIBUTORS.md +++ b/CONTRIBUTORS.md @@ -251,6 +251,7 @@ List of Contributors * [Piljae Chae](https://github.com/IHateMint) * [Oliver Kowalke](https://github.com/olk) * [Connor Goggins](https://github.com/connorgoggins) +* [Joe Evans](https://github.com/josephevans) Label Bot --------- diff --git a/src/operator/tensor/la_op-inl.h b/src/operator/tensor/la_op-inl.h index d580cced4ec5..7a5a602425fe 100644 --- a/src/operator/tensor/la_op-inl.h +++ b/src/operator/tensor/la_op-inl.h @@ -36,9 +36,10 @@ using namespace mshadow; // Copies lower/upper triangular part to upper/lower, i.e. to the opposite side. struct CopyTriangularToOppositeSide { template - MSHADOW_XINLINE static void Map(int i, int matrix_size, int stride, DType* data, bool to_lower) { + MSHADOW_XINLINE static void Map(index_t i, size_t matrix_size, index_t stride, + DType* data, bool to_lower) { // Below computation works even when we are dealing with a batch of matrices. - const int row((i % matrix_size) / stride), col(i % stride); + const index_t row((i % matrix_size) / stride), col(i % stride); if (row > col) { if (to_lower) { data[i] = data[i + (col - row) * (stride - 1)]; @@ -52,9 +53,9 @@ struct CopyTriangularToOppositeSide { // Zero's lower/upper triangular part of a matrix. struct ZeroTriangular { template - MSHADOW_XINLINE static void Map(int i, int matrix_size, int stride, DType* data, - bool zero_lower) { - const int row((i % matrix_size) / stride), col(i % stride); + MSHADOW_XINLINE static void Map(index_t i, size_t matrix_size, index_t stride, + DType* data, bool zero_lower) { + const index_t row((i % matrix_size) / stride), col(i % stride); if ((!zero_lower && (row < col)) || (zero_lower && (row > col))) data[i] = 0; } }; diff --git a/tests/nightly/test_large_array.py b/tests/nightly/test_large_array.py index d50f3ff7e539..4925e54c247b 100644 --- a/tests/nightly/test_large_array.py +++ b/tests/nightly/test_large_array.py @@ -37,6 +37,7 @@ LARGE_X = 100000000 SMALL_X = 100 SMALL_Y = 50 +LARGE_SQ_X = 80000 LARGE_SIZE = LARGE_X * SMALL_Y LARGE_TENSOR_SHAPE = 2**32 RNN_LARGE_TENSOR = 2**28 @@ -1167,6 +1168,32 @@ def check_correctness(mxnet_op, numpy_op, atol=1e-3): check_gather() check_binary_broadcast() +def test_linalg(): + def check_potrf(): + # creating an identity matrix input + A = nd.zeros((LARGE_SQ_X, LARGE_SQ_X)) + for i in range(LARGE_SQ_X): + A[i,i] = 1 + + out = nd.linalg.potrf(A) + # output should be an identity matrix + for i in range(LARGE_SQ_X): + assert out[i,i] == 1 + + def check_potri(): + # creating an identity matrix input + A = nd.zeros((LARGE_SQ_X, LARGE_SQ_X)) + for i in range(LARGE_SQ_X): + A[i,i] = 1 + + out = nd.linalg.potri(A) + # output should be an identity matrix + for i in range(LARGE_SQ_X): + assert out[i,i] == 1 + + check_potrf() + check_potri() + def test_basic(): def check_elementwise():