Cherry-pick large tensor support from apache#18752. (apache#18804)

Co-authored-by: Joe Evans <joeev@amazon.com>
zheyuye · Jul 29, 2020 · e9829e7 · e9829e7
1 parent 126636c
commit e9829e7
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md
@@ -254,6 +254,7 @@ List of Contributors
 * [Connor Goggins](https://github.com/connorgoggins)
 * [Wei Chu](https://github.com/waytrue17)
 * [Yang Shi](https://github.com/ys2843)
+* [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
@@ -36,9 +36,10 @@ using namespace mshadow;
 // Copies lower/upper triangular part to upper/lower, i.e. to the opposite side.
 struct CopyTriangularToOppositeSide {
   template<typename DType>
-  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<typename DType>
-  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;
   }
 };