diff --git a/src/operator/nn/moments-inl.h b/src/operator/nn/moments-inl.h index 6a9bdc54b905..ca78b65bf1ec 100644 --- a/src/operator/nn/moments-inl.h +++ b/src/operator/nn/moments-inl.h @@ -77,18 +77,18 @@ inline bool MomentsType(const nnvm::NodeAttrs& attrs, struct VarBroadcastKernel { template - 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; diff --git a/tests/nightly/test_np_large_array.py b/tests/nightly/test_np_large_array.py index c158977abfa4..1077d2d0c41d 100644 --- a/tests/nightly/test_np_large_array.py +++ b/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))