update module API for other submodules

python/mxnet/module/base_module.py

@@ -849,7 +849,7 @@ def get_input_grads(self, merge_multi_context=True):
         """
         raise NotImplementedError()
 
-    def update(self):
+    def update(self, sparse_pull_dict=None):
         """Updates parameters according to the installed optimizer and the gradients computed
         in the previous forward-backward batch.
 

python/mxnet/module/bucketing_module.py

@@ -334,7 +334,7 @@ def switch_bucket(self, bucket_key, data_shapes, label_shapes=None):
 
     def init_optimizer(self, kvstore='local', optimizer='sgd',
                        optimizer_params=(('learning_rate', 0.01),),
-                       force_init=False):
+                       force_init=False, sparse_pull_dict=None):
         """Installs and initializes optimizers.
 
         Parameters
@@ -356,7 +356,8 @@ def init_optimizer(self, kvstore='local', optimizer='sgd',
             return
 
         self._curr_module.init_optimizer(kvstore, optimizer, optimizer_params,
-                                         force_init=force_init)
+                                         force_init=force_init,
+                                         sparse_pull_dict=sparse_pull_dict)
         for mod in self._buckets.values():
             if mod is not self._curr_module:
                 mod.borrow_optimizer(self._curr_module)
@@ -399,13 +400,13 @@ def backward(self, out_grads=None):
         assert self.binded and self.params_initialized
         self._curr_module.backward(out_grads=out_grads)
 
-    def update(self):
+    def update(self, sparse_pull_dict=None):
         """Updates parameters according to installed optimizer and the gradient computed
         in the previous forward-backward cycle.
         """
         assert self.binded and self.params_initialized and self.optimizer_initialized
         self._params_dirty = True
-        self._curr_module.update()
+        self._curr_module.update(sparse_pull_dict=sparse_pull_dict)
 
     def get_outputs(self, merge_multi_context=True):
         """Gets outputs from a previous forward computation.

python/mxnet/module/python_module.py

@@ -110,7 +110,7 @@ def init_params(self, initializer=Uniform(0.01), arg_params=None, aux_params=Non
         """
         pass
 
-    def update(self):
+    def update(self, sparse_pull_dict=None):
         """Updates parameters according to the installed optimizer and the gradients computed
         in the previous forward-backward batch. Currently we do nothing here. Subclass should
         override this method if contains parameters.
@@ -196,7 +196,8 @@ def _compute_output_shapes(self):
         raise NotImplementedError()
 
     def init_optimizer(self, kvstore='local', optimizer='sgd',
-                       optimizer_params=(('learning_rate', 0.01),), force_init=False):
+                       optimizer_params=(('learning_rate', 0.01),), force_init=False,
+                       sparse_pull_dict=None):
         """Installs and initializes optimizers. By default we do nothing. Subclass should
         override this method if needed.
 

python/mxnet/module/sequential_module.py

@@ -275,7 +275,8 @@ def bind(self, data_shapes, label_shapes=None, for_training=True,
 
     def init_optimizer(self, kvstore='local', optimizer='sgd',
                        optimizer_params=(('learning_rate', 0.01),),
-                       force_init=False):
+                       force_init=False,
+                       sparse_pull_dict=None):
         """Installs and initializes optimizers.
 
         Parameters
@@ -298,7 +299,8 @@ def init_optimizer(self, kvstore='local', optimizer='sgd',
 
         for module in self._modules:
             module.init_optimizer(kvstore=kvstore, optimizer=optimizer,
-                                  optimizer_params=optimizer_params, force_init=force_init)
+                                  optimizer_params=optimizer_params, force_init=force_init,
+                                  sparse_pull_dict=sparse_pull_dict)
 
         self.optimizer_initialized = True
 
@@ -344,14 +346,14 @@ def backward(self, out_grads=None):
 
             out_grads = module.get_input_grads()
 
-    def update(self):
+    def update(self, sparse_pull_dict=None):
         """Updates parameters according to installed optimizer and the gradient computed
         in the previous forward-backward cycle.
         """
         assert self.binded and self.params_initialized and self.optimizer_initialized
 
         for module in self._modules:
-            module.update()
+            module.update(sparse_pull_dict=None)
 
     def get_outputs(self, merge_multi_context=True):
         """Gets outputs from a previous forward computation.