Numpy std var large tensor fix

src/operator/nn/moments-inl.h

@@ -77,18 +77,18 @@ inline bool MomentsType(const nnvm::NodeAttrs& attrs,
 
 struct VarBroadcastKernel {
   template<typename DType>
-  MSHADOW_XINLINE static void Map(int i,
+  MSHADOW_XINLINE static void Map(index_t i,
                                   DType *out,
                                   const DType *data,
                                   const DType *mean,
                                   mshadow::Shape<6> data_shape,
                                   mshadow::Shape<6> mean_shape) {
-    size_t data_idx = i;
-    size_t mean_idx = i;
-    size_t data_stride = 1;
-    size_t mean_stride = 1;
+    index_t data_idx = i;
+    index_t mean_idx = i;
+    index_t data_stride = 1;
+    index_t mean_stride = 1;
     for (int axis = 5; axis >= 0; --axis) {
-      size_t axis_idx = data_idx % data_shape[axis];
+      index_t axis_idx = data_idx % data_shape[axis];
       mean_idx -= axis_idx * data_stride;
       if (mean_shape[axis] != 1) {
         mean_idx += axis_idx * mean_stride;

tests/nightly/test_np_large_array.py

@@ -1256,7 +1256,7 @@ def test_diagflat():
     assert inp.grad.shape == inp.shape
     assert inp.grad[-1, -1] == 1
 
-    
+
 @use_np
 def test_diagonal():
     inp = np.zeros((2, INT_OVERFLOW+2))
@@ -1967,6 +1967,38 @@ def test_array_split():
     assert out[1][-1][-1] == 2
 
 
+@use_np
+def test_std():
+    N = 2*20
+    inp = np.zeros((2, INT_OVERFLOW))
+    inp[-1, -1] = N
+    inp.attach_grad()
+    with mx.autograd.record():
+        out = np.std(inp, axis=1)
+        out.backward()
+    assert out.shape == (2, )
+    ref = ((float(N)/INT_OVERFLOW)**2 * (INT_OVERFLOW-1))**0.5
+    assert_almost_equal(out[1], ref, rtol=1e-5, atol=1e-5)
+    assert inp.grad.shape == inp.shape
+    assert inp.grad[-1, -1] == 0
+
+
+@use_np
+def test_var():
+    N = 2*20
+    inp = np.zeros((2, INT_OVERFLOW))
+    inp[-1, -1] = N
+    inp.attach_grad()
+    with mx.autograd.record():
+        out = np.var(inp, axis=1)
+        out.backward()
+    assert out.shape == (2, )
+    ref = (float(N)/INT_OVERFLOW)**2 * (INT_OVERFLOW-1)
+    assert_almost_equal(out[1], ref, rtol=1e-5, atol=1e-5)
+
+    assert inp.grad.shape == inp.shape
+    assert inp.grad[-1, -1] == 0
+
 @use_np
 def test_rollaxis():
     inp = np.zeros((1, 1, 2, INT_OVERFLOW, 1))