[Numpy] Fix collect_params().zero_grad() in gluon numpy interface

python/mxnet/gluon/parameter.py

@@ -27,7 +27,6 @@
 from collections import OrderedDict, defaultdict
 import warnings
 import numpy as np
-import mxnet as mx
 
 from ..base import mx_real_t, MXNetError
 from .. import symbol, ndarray, initializer, context
@@ -896,15 +895,20 @@ def zero_grad(self):
                 continue
             for g in p.list_grad():
                 if g.stype == 'row_sparse':
-                    mx.ndarray.zeros_like(g, out=g)
+                    ndarray.zeros_like(g, out=g)
                 else:
                     arrays[g.context].append(g)
 
         if len(arrays) == 0:
             return
 
-        for arr in arrays.values():
-            mx.nd.reset_arrays(*arr, num_arrays=len(arr))
+        if is_np_array():
+            for arr in arrays.values():
+                for ele in arr:
+                    ele[()] = 0
+        else:
+            for arr in arrays.values():
+                ndarray.reset_arrays(*arr, num_arrays=len(arr))
 
     def reset_ctx(self, ctx):
         """Re-assign all Parameters to other contexts.

tests/python/unittest/test_numpy_gluon.py

@@ -179,6 +179,23 @@ def hybrid_forward(self, F, x, weight):
         assert_almost_equal(out.asnumpy(), (x.asnumpy() + const_arr), atol=1e-5, rtol=1e-4, use_broadcast=False)
 
 
+@use_np
+def test_parameters_zero_grad():
+    for hybridize in [False, True]:
+        net = gluon.nn.HybridSequential()
+        for _ in range(5):
+            net.add(gluon.nn.Dense(10))
+        if hybridize:
+            net.hybridize()
+        net.initialize()
+        out = net(mx.np.ones((32, 8)))
+        for v in net.collect_params().values():
+            v.grad()[()] = 1
+        net.collect_params().zero_grad()
+        for v in net.collect_params().values():
+            assert_almost_equal(v.grad().asnumpy(), mx.np.zeros_like(v.grad()).asnumpy())
+
+
 if __name__ == '__main__':
     import nose
     nose.runmodule()