Aggregated adamw update

python/mxnet/ndarray/contrib.py

@@ -32,6 +32,9 @@
 
 __all__ = ["rand_zipfian", "foreach", "while_loop", "cond", "isinf", "isfinite", "isnan"]
 
+def _flatten_list(nested_list):
+    return [item for sublist in nested_list for item in sublist]
+
 # pylint: disable=line-too-long
 def rand_zipfian(true_classes, num_sampled, range_max, ctx=None):
     """Draw random samples from an approximately log-uniform or Zipfian distribution.
@@ -514,7 +517,7 @@ def isfinite(data):
     [0. 0. 0. 1.]
     <NDArray 4 @cpu(0)>
     """
-    is_data_not_nan = data == data # pylint: disable=comparison-with-itself
+    is_data_not_nan = data == data  # pylint: disable=comparison-with-itself
     is_data_not_infinite = data.abs() != np.inf
     return ndarray.logical_and(is_data_not_infinite, is_data_not_nan)
 
@@ -542,14 +545,17 @@ def isnan(data):
     [1. 0.]
     <NDArray 2 @cpu(0)>
     """
-    return data != data # pylint: disable=comparison-with-itself
+    return data != data  # pylint: disable=comparison-with-itself
 
-def adamw_update(weight, grad, mean, var, rescale_grad, lr, eta, beta1=0.9, beta2=0.999,
-                 epsilon=1e-8, wd=0, clip_gradient=-1, out=None, name=None, **kwargs):
+def getRescaleGrad(rescale_grad, ctx):
     if not isinstance(rescale_grad, ndarray.NDArray):
-        rescale_grad = ndarray.full(shape=(1,), val=rescale_grad, ctx=weight.context)
+        return ndarray.full(shape=(1,), val=rescale_grad, ctx=ctx)
     else:
-        rescale_grad = rescale_grad.as_in_context(weight.context)
+        return rescale_grad.as_in_context(ctx)
+
+def adamw_update(weight, grad, mean, var, rescale_grad, lr, eta, beta1=0.9, beta2=0.999,
+                 epsilon=1e-8, wd=0, clip_gradient=-1, out=None, name=None, **kwargs):
+    rescale_grad = getRescaleGrad(rescale_grad, weight.context)
     return ndarray._internal._adamw_update(weight=weight, grad=grad, mean=mean, var=var,
                                            rescale_grad=rescale_grad, lr=lr, eta=eta,
                                            beta1=beta1, beta2=beta2, epsilon=epsilon,
@@ -559,13 +565,39 @@ def adamw_update(weight, grad, mean, var, rescale_grad, lr, eta, beta1=0.9, beta
 def mp_adamw_update(weight, grad, mean, var, weight32, rescale_grad, lr, eta, beta1=0.9,
                     beta2=0.999, epsilon=1e-8, wd=0, clip_gradient=-1, out=None,
                     name=None, **kwargs):
-    if not isinstance(rescale_grad, ndarray.NDArray):
-        rescale_grad = ndarray.full(shape=(1,), val=rescale_grad, ctx=weight.context)
-    else:
-        rescale_grad = rescale_grad.as_in_context(weight.context)
+    rescale_grad = getRescaleGrad(rescale_grad, weight.context)
     return ndarray._internal._mp_adamw_update(weight=weight, grad=grad, mean=mean, var=var,
                                               weight32=weight32,
                                               rescale_grad=rescale_grad, lr=lr, eta=eta,
                                               beta1=beta1, beta2=beta2, epsilon=epsilon,
                                               wd=wd, clip_gradient=clip_gradient, out=out,
                                               name=name, **kwargs)
+
+def multi_adamw_update(weights, grads, mean, var, rescale_grad, lrs, wds, etas,
+                       out=None, name=None, size=0, **kwargs):
+    if not size:
+        size = len(weights)
+    return ndarray._internal._multi_adamw_update(*_flatten_list(zip(weights, grads, mean, var)),
+                                                 getRescaleGrad(rescale_grad, weights[0].context),
+                                                 out=out,
+                                                 num_weights=size,
+                                                 lrs=lrs,
+                                                 wds=wds,
+                                                 etas=etas,
+                                                 name=name,
+                                                 **kwargs)
+
+def multi_mp_adamw_update(weights, grads, mean, var, weights32, rescale_grad, lrs, wds, etas,
+                          out=None, name=None, size=0, **kwargs):
+    if not size:
+        size = len(weights)
+    return ndarray._internal._multi_mp_adamw_update(
+        *_flatten_list(zip(weights, grads, mean, var, weights32)),
+        getRescaleGrad(rescale_grad, weights[0].context),
+        out=out,
+        num_weights=size,
+        lrs=lrs,
+        wds=wds,
+        etas=etas,
+        name=name,
+        **kwargs)