Allow large array operation in MXNet (dmlc#348)

azai91 · Jul 31, 2018 · d68d369 · d68d369
1 parent 463c0df
commit d68d369
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/mshadow/base.h b/mshadow/base.h
@@ -285,7 +285,7 @@ const unsigned kRandBufferSize = 1000000;
 /*! \brief pi  */
 const float kPi = 3.1415926f;
 /*! \brief type that will be used for index */
-typedef unsigned index_t;
+typedef int64_t index_t;
 
 #ifdef _WIN32
   /*! \brief openmp index for windows */

diff --git a/mshadow/tensor.h b/mshadow/tensor.h
@@ -123,8 +123,8 @@ struct Shape {
     return s;
   }
   /*! \return number of valid elements */
-  MSHADOW_XINLINE size_t Size(void) const {
-    size_t size = this->shape_[0];
+  MSHADOW_XINLINE index_t Size(void) const {
+    index_t size = this->shape_[0];
     #pragma unroll
     for (int i = 1; i < kDimension; ++i) {
       size *= this->shape_[i];
@@ -458,8 +458,8 @@ struct Tensor: public TRValue<Tensor<Device, dimension, DType>,
    * \tparam startdim the starting dimension
    */
   template<int startdim>
-  MSHADOW_XINLINE size_t MemSize(void) const {
-    size_t memsz = this->stride_;
+  MSHADOW_XINLINE index_t MemSize(void) const {
+    index_t memsz = this->stride_;
     #pragma unroll
     for (int i = startdim; i < kSubdim; ++i) {
       memsz *= this->shape_[i];
@@ -476,7 +476,7 @@ struct Tensor: public TRValue<Tensor<Device, dimension, DType>,
   /*!
    * \return memory cost of the tensor, including the aligned x dimension
    */
-  MSHADOW_XINLINE size_t MSize(void) const {
+  MSHADOW_XINLINE index_t MSize(void) const {
     return this->MemSize<0>();
   }
   /*!
@@ -582,7 +582,7 @@ struct Tensor<Device, 1, DType>:
   MSHADOW_XINLINE bool CheckContiguous(void) const {
     return true;
   }
-  MSHADOW_XINLINE size_t MSize(void) const {
+  MSHADOW_XINLINE index_t MSize(void) const {
     return shape_[0];
   }
   MSHADOW_XINLINE index_t size(index_t i) const {