From 3f2b68e934c5a756d07ae0fe7158014c7e85ab60 Mon Sep 17 00:00:00 2001
From: andreii <andreii@nvidia.com>
Date: Wed, 22 May 2019 09:25:35 -0700
Subject: [PATCH 1/9] Trigger CI


From 7af1a08a38f94485cd9cf535b2d1633b344bf0f4 Mon Sep 17 00:00:00 2001
From: andreii <andreii@nvidia.com>
Date: Tue, 8 Oct 2019 13:59:18 -0700
Subject: [PATCH 2/9] MxNet operator for aggregated Adam update

---
 python/mxnet/ndarray/contrib.py               |  52 ++-
 src/operator/contrib/adamw-inl.h              | 365 ++++++++++++++++--
 src/operator/contrib/adamw.cc                 | 166 ++++++--
 src/operator/contrib/adamw.cu                 |  34 +-
 tests/python/gpu/test_operator_gpu.py         |   1 +
 .../python/unittest/test_contrib_optimizer.py | 236 ++++++-----
 6 files changed, 664 insertions(+), 190 deletions(-)

diff --git a/python/mxnet/ndarray/contrib.py b/python/mxnet/ndarray/contrib.py
index 601bc682db38..d0eac246f4bf 100644
--- a/python/mxnet/ndarray/contrib.py
+++ b/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)
diff --git a/src/operator/contrib/adamw-inl.h b/src/operator/contrib/adamw-inl.h
index 6ae9e46b7def..3ff7bf4d0d32 100644
--- a/src/operator/contrib/adamw-inl.h
+++ b/src/operator/contrib/adamw-inl.h
@@ -18,26 +18,17 @@
  */
 
 /*!
- *  Copyright (c) 2016 by Contributors
- * \file optimizer_op-inl.h
+ *  Copyright (c) 2018 by Contributors
+ * \file adamw-inl.h
  * \brief Optimizer operators
- * \author Haibin Lin
+ * \author Haibin Lin, Moises Hernandez, Andrei Ivanov
  */
 #ifndef MXNET_OPERATOR_CONTRIB_ADAMW_INL_H_
 #define MXNET_OPERATOR_CONTRIB_ADAMW_INL_H_
-#include <dmlc/parameter.h>
 #include <mxnet/operator.h>
-#include <mxnet/operator_util.h>
-#include <mxnet/op_attr_types.h>
-#include <mshadow/base.h>
-#include <nnvm/op.h>
-#include <nnvm/op_attr_types.h>
 #include <vector>
-#include <cmath>
-#include "../operator_common.h"
 #include "../mshadow_op.h"
 #include "../elemwise_op_common.h"
-#include "../mxnet_op.h"
 
 namespace mxnet {
 namespace op {
@@ -120,20 +111,14 @@ struct MPAdamWKernel {
     const float param_eta, const float param_lr, const float param_wd,
     const float param_rescale_grad, const float param_epsilon) {
     float w = weight32[i];
-    float mean = mean_data[i];
-    float var = var_data[i];
     float scaled_grad = param_rescale_grad*static_cast<float>(grad_data[i]);
-    if (param_clip_gradient >= 0.0f) {
-      mean = param_beta1 * mean +
-             (1 - param_beta1) * mshadow_op::clip::Map(scaled_grad, param_clip_gradient);
-      var = param_beta2 * var + (1 - param_beta2) *
-            mshadow_op::square::Map(mshadow_op::clip::Map(scaled_grad, param_clip_gradient));
-    } else {
-      mean = param_beta1 * mean + (1 - param_beta1) * scaled_grad;
-      var = param_beta2 * var + (1 - param_beta2) * mshadow_op::square::Map(scaled_grad);
-    }
-    mean_data[i] = mean;
-    var_data[i] = var;
+    if (param_clip_gradient >= 0.0f)
+      scaled_grad = mshadow_op::clip::Map(scaled_grad, param_clip_gradient);
+
+    float mean = mean_data[i] = param_beta1 * mean_data[i] + (1.0f - param_beta1) * scaled_grad;
+    float var = var_data[i] = param_beta2 * var_data[i] +
+                  (1.0f - param_beta2) * mshadow_op::square::Map(scaled_grad);
+
     w = w - param_eta * (param_lr * mean / (mshadow_op::square_root::Map(var) + param_epsilon)
                          + param_wd * w);
     weight32[i] = w;
@@ -141,7 +126,6 @@ struct MPAdamWKernel {
   }
 };
 
-
 template<typename xpu>
 struct MPAdamWUpdate {
   static inline void Forward(const nnvm::NodeAttrs& attrs,
@@ -151,7 +135,7 @@ struct MPAdamWUpdate {
                const std::vector<TBlob> &outputs,
                const float rescale_grad) {
     using namespace mxnet_op;
-    AdamWParam param = nnvm::get<AdamWParam>(attrs.parsed);
+    const auto &param = nnvm::get<AdamWParam>(attrs.parsed);
     Stream<xpu>* s = ctx.get_stream<xpu>();
     MSHADOW_REAL_TYPE_SWITCH(inputs[0].type_flag_, DType, {
       Tensor<xpu, 2, DType> weight = inputs[0].FlatTo2D<xpu, DType>(s);
@@ -183,25 +167,22 @@ struct AdamWUpdate {
     using namespace mshadow;
     using namespace mshadow::expr;
     using namespace mshadow_op;
-    const AdamWParam& param = nnvm::get<AdamWParam>(attrs.parsed);
+    const auto &param = nnvm::get<AdamWParam>(attrs.parsed);
     Stream<xpu>* s = ctx.get_stream<xpu>();
     MSHADOW_REAL_TYPE_SWITCH(inputs[0].type_flag_, DType, {
-      Tensor<xpu, 2, DType> weight = inputs[0].FlatTo2D<xpu, DType>(s);
+      const Tensor<xpu, 2, DType> &weight = inputs[0].FlatTo2D<xpu, DType>(s);
       Tensor<xpu, 2, DType> grad = inputs[1].FlatTo2D<xpu, DType>(s);
       Tensor<xpu, 2, DType> mean = inputs[2].FlatTo2D<xpu, DType>(s);
       Tensor<xpu, 2, DType> var = inputs[3].FlatTo2D<xpu, DType>(s);
       Tensor<xpu, 2, DType> out = outputs[0].FlatTo2D<xpu, DType>(s);
 
       grad = scalar<DType>(rescale_grad) * grad;
-      if (param.clip_gradient >= 0.0f) {
-        mean = scalar<DType>(param.beta1)*mean + scalar<DType>(1.f-param.beta1) *
-            F<clip>(grad, DType(param.clip_gradient));
-        var = scalar<DType>(param.beta2)*var + scalar<DType>(1.f-param.beta2)*F<square>(
-            F<clip>(grad, DType(param.clip_gradient)));
-      } else {
-        mean = scalar<DType>(param.beta1)*mean + scalar<DType>(1.f-param.beta1) * grad;
-        var = scalar<DType>(param.beta2)*var + scalar<DType>(1.f-param.beta2) * F<square>(grad);
-      }
+      if (param.clip_gradient >= 0.0f)
+        grad = F<clip>(grad, DType(param.clip_gradient));
+
+      mean = scalar<DType>(param.beta1) * mean + scalar<DType>(1.f-param.beta1) * grad;
+      var = scalar<DType>(param.beta2) * var + scalar<DType>(1.f-param.beta2) * F<square>(grad);
+
       Assign(out, req[0],
              weight -
              scalar<DType>(param.eta) * (scalar<DType>(param.lr) *
@@ -211,6 +192,314 @@ struct AdamWUpdate {
   }
 };
 
+////
+// Multiple gradients in single kernel
+////
+
+struct MultiAdamWParam : public dmlc::Parameter<MultiAdamWParam> {
+  mxnet::Tuple<float> lrs;
+  mxnet::Tuple<float> wds;
+  mxnet::Tuple<float> etas;
+  float beta1;
+  float beta2;
+  float epsilon;
+  float rescale_grad;
+  float clip_gradient;
+  int num_weights;
+  DMLC_DECLARE_PARAMETER(MultiAdamWParam) {
+    DMLC_DECLARE_FIELD(lrs)
+    .describe("Learning rates");
+    DMLC_DECLARE_FIELD(beta1)
+    .set_default(0.9f)
+    .describe("The decay rate for the 1st moment estimates.");
+    DMLC_DECLARE_FIELD(beta2)
+    .set_default(0.999f)
+    .describe("The decay rate for the 2nd moment estimates.");
+    DMLC_DECLARE_FIELD(epsilon)
+    .set_default(1e-8f)
+    .describe("A small constant for numerical stability.");
+    DMLC_DECLARE_FIELD(wds)
+    .describe("Weight decay augments the objective function with a "
+              "regularization term that penalizes large weights. "
+              "The penalty scales with the square of the magnitude of each weight.");
+    DMLC_DECLARE_FIELD(etas)
+    .describe("Learning rates schedule multiplier");
+    DMLC_DECLARE_FIELD(clip_gradient)
+    .set_default(-1.0f)
+    .describe("Clip gradient to the range of [-clip_gradient, clip_gradient] "
+              "If clip_gradient <= 0, gradient clipping is turned off. "
+              "grad = max(min(grad, clip_gradient), -clip_gradient).");
+    DMLC_DECLARE_FIELD(num_weights)
+    .set_default(1)
+    .describe("Number of updated weights.");
+  }
+};
+
+
+template<typename ParamType, int input_stride>
+inline bool MP_MultiAdamW_InferShape(const nnvm::NodeAttrs& attrs,
+                                          mxnet::ShapeVector *in_attrs,
+                                          mxnet::ShapeVector *out_attrs) {
+  const ParamType& param = dmlc::get<ParamType>(attrs.parsed);
+  CHECK_EQ(in_attrs->size(), input_stride * param.num_weights +1);
+  CHECK_EQ(out_attrs->size(), param.num_weights);
+
+  bool all_inferred = true;
+  auto& input_shapes = *in_attrs;
+  auto& output_shapes = *out_attrs;
+
+  // Learning rates
+  CHECK_EQ(param.lrs.ndim(), param.num_weights)
+    << "Number of learning rates is inconsistent with num_weights "
+    << "parameter passed. Expected number of learning rates: "
+    << param.num_weights << ", and got " << param.lrs.ndim();
+  // Weight decays
+  CHECK_EQ(param.wds.ndim(), param.num_weights)
+    << "Number of weight decays is inconsistent with num_weights "
+    << "parameter passed. Expected number of weight decays: "
+    << param.num_weights << ", and got " << param.wds.ndim();
+  // Learning rates schedule multiplier
+  CHECK_EQ(param.etas.ndim(), param.num_weights)
+    << "Number of learning rates schedule multiplier is inconsistent with num_weights "
+    << "parameter passed. Expected number of learning rates schedule multiplier: "
+    << param.num_weights << ", and got " << param.lrs.ndim();
+
+  // Weights, gradients, mean and variance
+  for (int i = 0; i < param.num_weights; ++i) {
+    mxnet::ShapeVector input_vec;
+    mxnet::ShapeVector output_vec({output_shapes[i]});
+    for (int j = 0; j < input_stride; ++j) {
+      input_vec.push_back(input_shapes[i * input_stride + j]);
+    }
+    all_inferred = all_inferred && ElemwiseShape<input_stride, 1>(attrs, &input_vec, &output_vec);
+  }
+  // rescale_grad.shape = ()
+  SHAPE_ASSIGN_CHECK(*in_attrs, param.num_weights*input_stride, mxnet::TShape());
+  return all_inferred;
+}
+
+template <typename ParamType, int input_stride, int num_fp32_inputs>
+inline bool MP_MultiAdamW_InferType(const nnvm::NodeAttrs& attrs,
+                                    std::vector<int> *in_attrs,
+                                    std::vector<int> *out_attrs) {
+  const ParamType& param = dmlc::get<ParamType>(attrs.parsed);
+  CHECK_EQ(in_attrs->size(), input_stride * param.num_weights +1);
+  CHECK_EQ(out_attrs->size(), param.num_weights);
+
+  bool all_inferred = true;
+  auto& input_types = *in_attrs;
+  auto& output_types = *out_attrs;
+
+  // Weights, gradients,
+  for (int i = 0; i < param.num_weights; ++i) {
+    std::vector<int> input_vec;
+    std::vector<int> output_vec({output_types[i]});
+    for (int j = 0; j < input_stride - 2 - num_fp32_inputs; ++j) {
+      input_vec.push_back(input_types[i * input_stride + j]);
+    }
+    all_inferred = all_inferred &&
+            ElemwiseType<input_stride - 2 - num_fp32_inputs, 1>(attrs, &input_vec, &output_vec);
+  }
+  // mean, var
+  for (int i = 0; i < param.num_weights; ++i) {
+    TYPE_ASSIGN_CHECK(input_types, input_stride * i +2, mshadow::kFloat32);
+    TYPE_ASSIGN_CHECK(input_types, input_stride * i +3, mshadow::kFloat32);
+  }
+
+  // master copies of weights
+  for (int i = 0; i < param.num_weights; ++i) {
+    for (int j = 0; j < num_fp32_inputs; ++j) {
+      TYPE_ASSIGN_CHECK(input_types, input_stride * i + input_stride - 1 - j, mshadow::kFloat32);
+    }
+  }
+  // rescale_grad.type = ()
+  TYPE_ASSIGN_CHECK(input_types, param.num_weights*input_stride, mshadow::kFloat32);
+  return all_inferred;
+}
+
+
+template<typename T>
+class Adam_type_identity {
+ public:
+  using type = T;
+};
+
+
+template<typename T>
+class Adam_single_precision {
+ public:
+  using type = float;
+};
+
+template<typename DType, typename MPDType>
+struct MultiAdamKernelParam {
+  static const int N = 50;
+  int count;
+  size_t max_size;
+  size_t sizes[N];
+  DType* weights[N];
+  DType* grad_data[N];
+  MPDType* mean_data[N];
+  MPDType* var_data[N];
+  MPDType* weights32[N];
+  DType* out_data[N];
+  MPDType clip_gradient;
+  MPDType beta1;
+  MPDType beta2;
+  MPDType etas[N];
+  MPDType lrs[N];
+  MPDType wds[N];
+  MPDType epsilon;
+};
+
+template<typename MPDType, bool has_mixed_precision>
+struct MultiMPAdamWKernel {
+  template<typename DType>
+  MSHADOW_XINLINE static void Map(int i, const MultiAdamKernelParam<DType, MPDType>& param,
+                                  const OpReqType req, const float rescale_grad){
+    for (int index = 0; index < param.count; ++index) {
+      if ((size_t)i < param.sizes[index]) {
+        MPDType w = has_mixed_precision ? param.weights32[index][i]:
+                                          MPDType(param.weights[index][i]);
+        MPDType scaled_grad = static_cast<MPDType>(rescale_grad)*
+                              static_cast<MPDType>(param.grad_data[index][i]);
+
+        if (param.clip_gradient >= 0.0f)
+          scaled_grad = mshadow_op::clip::Map(scaled_grad, param.clip_gradient);
+
+        const auto mean = param.beta1 * (param.mean_data[index][i]- scaled_grad) + scaled_grad;
+        const auto adj = mshadow_op::square::Map(scaled_grad);
+        const auto var = param.beta2 * (param.var_data[index][i] - adj) + adj;
+
+        param.mean_data[index][i] = mean;
+        param.var_data[index][i] = var;
+        w = w - param.etas[index] * (param.lrs[index] *
+            mean / (mshadow_op::square_root::Map(var) + param.epsilon)
+            + param.wds[index] * w);
+        if (has_mixed_precision)
+          param.weights32[index][i] = w;
+
+        KERNEL_ASSIGN(param.out_data[index][i], req, w);
+      }
+    }
+  }
+};
+
+template<typename xpu,
+         typename DType,
+         typename MPDType,
+         typename ParamType = MultiAdamWParam,
+         int input_stride = 4>
+void FillMultiAdamKernelParam(const nnvm::NodeAttrs& attrs,
+                              const OpContext &ctx,
+                              const std::vector<TBlob> &inputs,
+                              const std::vector<TBlob> &outputs,
+                              MultiAdamKernelParam<DType, MPDType> *pParam) {
+  const ParamType& p = nnvm::get<ParamType>(attrs.parsed);
+  mxnet_op::Stream<xpu>* s = ctx.get_stream<xpu>();
+  pParam->clip_gradient = p.clip_gradient;
+  pParam->beta1 = p.beta1;
+  pParam->beta2 = p.beta2;
+
+  pParam->epsilon = p.epsilon;
+
+  pParam->count = p.num_weights;
+  pParam->max_size = 0;
+  const bool isSame = std::is_same<DType, MPDType>::value;
+  for (int i = 0; i < pParam->count; ++i) {
+    const auto idx = i * input_stride;
+    pParam->sizes[i] = inputs[idx].shape_.Size();
+    if (pParam->max_size < pParam->sizes[i])
+      pParam->max_size = pParam->sizes[i];
+
+    pParam->weights[i] = inputs[idx].FlatTo2D<xpu, DType>(s).dptr_;
+    pParam->grad_data[i] = inputs[idx + 1].FlatTo2D<xpu, DType>(s).dptr_;
+    pParam->mean_data[i] = inputs[idx + 2].FlatTo2D<xpu, MPDType>(s).dptr_;
+    pParam->var_data[i]  = inputs[idx + 3].FlatTo2D<xpu, MPDType>(s).dptr_;
+    // if mixed precision, then the last input in a set
+    // is 32-bit master copy of the weights
+    if (!isSame)
+      pParam->weights32[i] = inputs[idx + input_stride - 1].FlatTo2D<xpu, MPDType>(s).dptr_;
+
+    pParam->out_data[i] = outputs[i].FlatTo2D<xpu, DType>(s).dptr_;
+  }
+  memcpy(pParam->etas, p.etas.begin(), pParam->count * sizeof(p.etas[0]));
+  memcpy(pParam->lrs, p.lrs.begin(), pParam->count * sizeof(p.lrs[0]));
+  memcpy(pParam->wds, p.wds.begin(), pParam->count * sizeof(p.wds[0]));
+}
+
+template<typename xpu, template<typename> class MPTypeChooser, int input_stride>
+static inline void MultiAdamWUpdate(const nnvm::NodeAttrs& attrs,
+                                    const OpContext &ctx,
+                                    const std::vector<TBlob> &inputs,
+                                    const std::vector<OpReqType> &req,
+                                    const std::vector<TBlob> &outputs,
+                                    const float rescale_grad) {
+  using namespace mxnet_op;
+  Stream<xpu>* s = ctx.get_stream<xpu>();
+  MSHADOW_REAL_TYPE_SWITCH(outputs[0].type_flag_, DType, {
+    using MPDType = typename MPTypeChooser<DType>::type;
+    MultiAdamKernelParam<DType, MPDType> param;
+    FillMultiAdamKernelParam<xpu, DType, MPDType, MultiAdamWParam, input_stride>
+            (attrs, ctx, inputs, outputs, &param);
+
+    Kernel<MultiMPAdamWKernel<MPDType, !std::is_same<DType, MPDType>::value>, xpu>::
+                              Launch(s, param.max_size, param, req[0], rescale_grad);
+  });
+}
+
+template<typename xpu>
+void GetScaleFloat(const TBlob &scale_blob, float *pScalef);
+
+template<typename xpu>
+bool PrepareInputBlobs(const std::vector<TBlob> &inputs,
+                       std::vector<TBlob> *inputs_wo_scale,
+                       float *pScalef) {
+  const size_t num_in = inputs.size() - 1;
+  GetScaleFloat<xpu>(inputs[num_in], pScalef);
+  if (!std::isfinite(*pScalef) || *pScalef == 0)
+    return false;
+
+  inputs_wo_scale->reserve(num_in);
+  for (size_t i = 0; i < num_in; i++)
+    inputs_wo_scale->emplace_back(inputs[i]);
+
+  return true;
+}
+
+template<typename xpu, class F>
+inline void MPUpdate(const nnvm::NodeAttrs& attrs,
+                     const OpContext &ctx,
+                     const std::vector<TBlob> &inputs,
+                     const std::vector<OpReqType> &req,
+                     const std::vector<TBlob> &outputs) {
+  std::vector<TBlob> inputs_wo_scale;
+  float scalef;
+  if (!PrepareInputBlobs<xpu>(inputs, &inputs_wo_scale, &scalef))
+    return;
+
+  F::Forward(attrs, ctx, inputs_wo_scale, req, outputs, scalef);
+}
+
+template<typename xpu, bool MP>
+inline void multiMPUpdate(const nnvm::NodeAttrs& attrs,
+                          const OpContext &ctx,
+                          const std::vector<TBlob> &inputs,
+                          const std::vector<OpReqType> &req,
+                          const std::vector<TBlob> &outputs) {
+  std::vector<TBlob> inputs_wo_scale;
+  float scalef;
+  if (!PrepareInputBlobs<xpu>(inputs, &inputs_wo_scale, &scalef))
+    return;
+
+  if (!MP)
+    MultiAdamWUpdate<xpu, Adam_type_identity, 4>
+      (attrs, ctx, inputs_wo_scale, req, outputs, scalef);
+  else
+    MultiAdamWUpdate<xpu, Adam_single_precision, 5>
+      (attrs, ctx, inputs_wo_scale, req, outputs, scalef);
+}
+
 }  // namespace op
 }  // namespace mxnet
 
diff --git a/src/operator/contrib/adamw.cc b/src/operator/contrib/adamw.cc
index f0716c6020f9..2c730f0b3e7b 100644
--- a/src/operator/contrib/adamw.cc
+++ b/src/operator/contrib/adamw.cc
@@ -18,37 +18,18 @@
  */
 
 /*!
- *  Copyright (c) 2016 by Contributors
- * \file optimizer_op.cc
+ *  Copyright (c) 2018 by Contributors
+ * \file adamw.cc
  * \brief Optimizer operators
- * \author Haibin Lin
+ * \author Haibin Lin, Moises Hernandez, Andrei Ivanov
  */
 #include "./adamw-inl.h"
-#include "../optimizer_op-inl.h"
 
 namespace mxnet {
 namespace op {
 
 DMLC_REGISTER_PARAMETER(AdamWParam);
-
-template<template <typename xpu> class F>
-inline void MPUpdateCPU(const nnvm::NodeAttrs& attrs,
-                        const OpContext &ctx,
-                        const std::vector<TBlob> &inputs,
-                        const std::vector<OpReqType> &req,
-                        const std::vector<TBlob> &outputs) {
-  // copy to cpu and check NaN value
-  TBlob scale_blob = inputs[inputs.size() - 1];
-  MSHADOW_REAL_TYPE_SWITCH(scale_blob.type_flag_, DType, {
-    float scalef = static_cast<float>(*scale_blob.dptr<DType>());
-    if (!std::isfinite(scalef) || scalef == 0) return;
-    std::vector<TBlob> inputs_wo_scale;
-    size_t num_in = inputs.size();
-    inputs_wo_scale.reserve(num_in - 1);
-    for (size_t i = 0; i < num_in - 1; i++) inputs_wo_scale.emplace_back(inputs[i]);
-    F<cpu>::Forward(attrs, ctx, inputs_wo_scale, req, outputs, scalef);
-  });
-}
+DMLC_REGISTER_PARAMETER(MultiAdamWParam);
 
 NNVM_REGISTER_OP(_mp_adamw_update)
 .describe(R"code(Update function for multi-precision AdamW optimizer.
@@ -84,7 +65,7 @@ the update is skipped.
   [](const nnvm::NodeAttrs& attrs) {
     return std::vector<uint32_t>{2, 3, 4};
   })
-.set_attr<FCompute>("FCompute<cpu>", MPUpdateCPU<MPAdamWUpdate>)
+.set_attr<FCompute>("FCompute<cpu>", MPUpdate<cpu, MPAdamWUpdate<cpu>>)
 .add_argument("weight", "NDArray-or-Symbol", "Weight")
 .add_argument("grad", "NDArray-or-Symbol", "Gradient")
 .add_argument("mean", "NDArray-or-Symbol", "Moving mean")
@@ -127,7 +108,7 @@ the update is skipped.
   [](const nnvm::NodeAttrs& attrs) {
     return std::vector<uint32_t>{2, 3};
   })
-.set_attr<FCompute>("FCompute<cpu>", MPUpdateCPU<AdamWUpdate>)
+.set_attr<FCompute>("FCompute<cpu>", MPUpdate<cpu, AdamWUpdate<cpu>>)
 .add_argument("weight", "NDArray-or-Symbol", "Weight")
 .add_argument("grad", "NDArray-or-Symbol", "Gradient")
 .add_argument("mean", "NDArray-or-Symbol", "Moving mean")
@@ -137,5 +118,140 @@ the update is skipped.
               "the update is skipped.")
 .add_arguments(AdamWParam::__FIELDS__());
 
+template<>
+void GetScaleFloat<cpu>(const TBlob &scale_blob, float *pScalef) {
+  MSHADOW_REAL_TYPE_SWITCH(scale_blob.type_flag_, DType,
+    *pScalef = static_cast<float>(*scale_blob.dptr<DType>());
+  )
+}
+
+std::vector<std::string> ParamToVector(uint32_t num_args, const char *pName[], size_t nParams) {
+  std::vector<std::string> ret;
+  for (uint32_t i = 0; i < num_args; ++i) {
+    const auto idx = std::to_string(i);
+    for (size_t j = 0; j < nParams; ++j)
+      ret.push_back(std::string(pName[i]) + idx);
+  }
+
+  return ret;
+}
+
+inline uint32_t num_weights(const nnvm::NodeAttrs& attrs) {
+  return static_cast<uint32_t>(dmlc::get<MultiAdamWParam>(attrs.parsed).num_weights);
+}
+
+NNVM_REGISTER_OP(_multi_adamw_update)
+.describe(R"code(Update function for AdamW optimizer.
+
+AdamW is seen as a modification of Adam by decoupling the weight decay from the
+optimization steps taken w.r.t. the loss function.
+
+Adam update consists of the following steps, where g represents gradient and m, v
+are 1st and 2nd order moment estimates (mean and variance).
+
+.. math::
+
+ g_t = \nabla J(W_{t-1})\\
+ m_t = \beta_1 m_{t-1} + (1 - \beta_1) g_t\\
+ v_t = \beta_2 v_{t-1} + (1 - \beta_2) g_t^2\\
+ W_t = W_{t-1} - \eta_t (\alpha \frac{ m_t }{ \sqrt{ v_t } + \epsilon } + wd W_{t-1})
+
+It updates the weights using::
+
+ m = beta1*m + (1-beta1)*grad
+ v = beta2*v + (1-beta2)*(grad**2)
+ w -= eta * (learning_rate * m / (sqrt(v) + epsilon) + w * wd)
+
+Note that gradient is rescaled to grad = rescale_grad * grad. If rescale_grad is NaN, Inf, or 0,
+the update is skipped.
+)code" ADD_FILELINE)
+.set_num_inputs([](const nnvm::NodeAttrs& attrs) {
+    return num_weights(attrs) * 4 + 1;
+  })
+.set_num_outputs([](const nnvm::NodeAttrs& attrs) {
+    return num_weights(attrs);
+  })
+.set_attr_parser(ParamParser<MultiAdamWParam>)
+.set_attr<mxnet::FInferShape>("FInferShape", MP_MultiAdamW_InferShape<MultiAdamWParam, 4>)
+.set_attr<nnvm::FInferType>("FInferType", ElemwiseType<-1, -1>)
+.set_attr<nnvm::FListInputNames>("FListInputNames",
+  [](const NodeAttrs& attrs) {
+    const char *paramName[] = {"weight_", "grad_", "mean_", "var_", "rescale_grad_"};
+    return ParamToVector(num_weights(attrs), paramName, sizeof(paramName)/sizeof(paramName[0]));
+  })
+// mutable: mean, var
+.set_attr<nnvm::FMutateInputs>("FMutateInputs",
+  [](const nnvm::NodeAttrs& attrs) {
+    std::vector<uint32_t> ret;
+    const auto iMax = num_weights(attrs);
+    for (size_t i = 0; i < iMax; ++i) {
+      ret.push_back(i * 4 + 2);
+      ret.push_back(i * 4 + 3);
+    }
+    return ret;
+  })
+
+.set_attr<FCompute>("FCompute<cpu>", multiMPUpdate<cpu, false>)
+.add_argument("data", "NDArray-or-Symbol[]", "data")
+.add_arguments(MultiAdamWParam::__FIELDS__());
+
+
+NNVM_REGISTER_OP(_multi_mp_adamw_update)
+.describe(R"code(Update function for multi-precision AdamW optimizer.
+
+AdamW is seen as a modification of Adam by decoupling the weight decay from the
+optimization steps taken w.r.t. the loss function.
+
+Adam update consists of the following steps, where g represents gradient and m, v
+are 1st and 2nd order moment estimates (mean and variance).
+
+.. math::
+
+ g_t = \nabla J(W_{t-1})\\
+ m_t = \beta_1 m_{t-1} + (1 - \beta_1) g_t\\
+ v_t = \beta_2 v_{t-1} + (1 - \beta_2) g_t^2\\
+ W_t = W_{t-1} - \eta_t (\alpha \frac{ m_t }{ \sqrt{ v_t } + \epsilon } + wd W_{t-1})
+
+It updates the weights using::
+
+ m = beta1*m + (1-beta1)*grad
+ v = beta2*v + (1-beta2)*(grad**2)
+ w -= eta * (learning_rate * m / (sqrt(v) + epsilon) + w * wd)
+
+Note that gradient is rescaled to grad = rescale_grad * grad. If rescale_grad is NaN, Inf, or 0,
+the update is skipped.
+)code" ADD_FILELINE)
+.set_num_inputs([](const nnvm::NodeAttrs& attrs) {
+    return num_weights(attrs) * 5 + 1;
+  })
+.set_num_outputs([](const nnvm::NodeAttrs& attrs) {
+    return num_weights(attrs);
+  })
+.set_attr_parser(ParamParser<MultiAdamWParam>)
+.set_attr<mxnet::FInferShape>("FInferShape", MP_MultiAdamW_InferShape<MultiAdamWParam, 5>)
+.set_attr<nnvm::FInferType>("FInferType", MP_MultiAdamW_InferType<MultiAdamWParam, 5, 1>)
+.set_attr<nnvm::FListInputNames>("FListInputNames",
+  [](const NodeAttrs& attrs) {
+    const char *paramName[] = {"weight_", "grad_", "mean_", "var_", "weight32_", "rescale_grad_"};
+    return ParamToVector(num_weights(attrs), paramName, sizeof(paramName)/sizeof(paramName[0]));
+  })
+// mutable: mean, var, weights32
+.set_attr<nnvm::FMutateInputs>("FMutateInputs",
+  [](const nnvm::NodeAttrs& attrs) {
+    std::vector<uint32_t> ret;
+    const auto iMax = num_weights(attrs);
+    for (size_t i = 0; i < iMax; ++i) {
+      ret.push_back(i * 5 + 2);
+      ret.push_back(i * 5 + 3);
+      ret.push_back(i * 5 + 4);
+    }
+    return ret;
+  })
+
+.set_attr<FCompute>("FCompute<cpu>", multiMPUpdate<cpu, true>)
+.add_argument("data", "NDArray-or-Symbol[]", "data")
+.add_arguments(MultiAdamWParam::__FIELDS__());
+
+
 }  // namespace op
 }  // namespace mxnet
diff --git a/src/operator/contrib/adamw.cu b/src/operator/contrib/adamw.cu
index 1521749904b9..81b13c98ede4 100644
--- a/src/operator/contrib/adamw.cu
+++ b/src/operator/contrib/adamw.cu
@@ -21,40 +21,34 @@
  *  Copyright (c) 2018 by Contributors
  * \file adamw.cu
  * \brief Optimizer operators
- * \author Haibin Lin
+ * \author Haibin Lin, Moises Hernandez, Andrei Ivanov
  */
 #include "./adamw-inl.h"
 
 namespace mxnet {
 namespace op {
 
-template<template <typename xpu> class F>
-inline void MPUpdateGPU(const nnvm::NodeAttrs& attrs,
-                        const OpContext &ctx,
-                        const std::vector<TBlob> &inputs,
-                        const std::vector<OpReqType> &req,
-                        const std::vector<TBlob> &outputs) {
-  // copy to cpu and check NaN value
-  TBlob scale_blob = inputs[inputs.size() - 1];
+template<>
+void GetScaleFloat<gpu>(const TBlob &scale_blob, float *pScalef) {
   MSHADOW_REAL_TYPE_SWITCH(scale_blob.type_flag_, DType, {
     DType scale = 0;
     CUDA_CALL(cudaMemcpy(&scale, scale_blob.dptr<DType>(), sizeof(DType),
-       cudaMemcpyDeviceToHost));
-    float scalef = static_cast<float>(scale);
-    if (!std::isfinite(scalef) || scalef == 0) return;
-    std::vector<TBlob> inputs_wo_scale;
-    size_t num_in = inputs.size();
-    inputs_wo_scale.reserve(num_in - 1);
-    for (size_t i = 0; i < num_in - 1; i++) inputs_wo_scale.emplace_back(inputs[i]);
-    F<gpu>::Forward(attrs, ctx, inputs_wo_scale, req, outputs, scalef);
-  });
+                         cudaMemcpyDeviceToHost));
+    *pScalef = static_cast<float>(scale);
+  })
 }
 
 NNVM_REGISTER_OP(_adamw_update)
-.set_attr<FCompute>("FCompute<gpu>", MPUpdateGPU<AdamWUpdate>);
+.set_attr<FCompute>("FCompute<gpu>", MPUpdate<gpu, AdamWUpdate<gpu>>);
 
 NNVM_REGISTER_OP(_mp_adamw_update)
-.set_attr<FCompute>("FCompute<gpu>", MPUpdateGPU<MPAdamWUpdate>);
+.set_attr<FCompute>("FCompute<gpu>", MPUpdate<gpu, MPAdamWUpdate<gpu>>);
+
+NNVM_REGISTER_OP(_multi_adamw_update)
+.set_attr<FCompute>("FCompute<gpu>", multiMPUpdate<gpu, false>);
+
+NNVM_REGISTER_OP(_multi_mp_adamw_update)
+.set_attr<FCompute>("FCompute<gpu>", multiMPUpdate<gpu, true>);
 
 }  // namespace op
 }  // namespace mxnet
diff --git a/tests/python/gpu/test_operator_gpu.py b/tests/python/gpu/test_operator_gpu.py
index e3862abbebd1..f2a67d828705 100644
--- a/tests/python/gpu/test_operator_gpu.py
+++ b/tests/python/gpu/test_operator_gpu.py
@@ -49,6 +49,7 @@
 from test_contrib_operator import test_multibox_target_op
 from test_tvm_op import *
 from test_library_loading import *
+from test_contrib_optimizer import test_adamw
 
 set_default_context(mx.gpu(0))
 del test_support_vector_machine_l1_svm  # noqa
diff --git a/tests/python/unittest/test_contrib_optimizer.py b/tests/python/unittest/test_contrib_optimizer.py
index 675cc94c64f1..d42999d00c28 100644
--- a/tests/python/unittest/test_contrib_optimizer.py
+++ b/tests/python/unittest/test_contrib_optimizer.py
@@ -16,12 +16,14 @@
 # under the License.
 
 import itertools
-
 import numpy as np
-
 import mxnet as mx
 from mxnet.test_utils import *
 
+curr_path = os.path.dirname(os.path.abspath(os.path.expanduser(__file__)))
+sys.path.insert(0, os.path.join(curr_path, '../unittest'))
+from common import with_seed
+
 
 # * GroupAdaGrad
 class PyGroupAdaGrad(mx.optimizer.Optimizer):
@@ -94,102 +96,142 @@ def test_group_adagrad():
                 g_stype='row_sparse',
                 compare_states=False)
 
-def test_adamw():
-    shape = (3, 4)
-    weight = mx.nd.random.uniform(shape=shape)
-    weight_ref = weight.copy()
-    grad = mx.nd.random.uniform(shape=shape)
-    m = mx.nd.random.uniform(shape=shape)
-    v = mx.nd.random.uniform(shape=shape)
-    rescale_grad = mx.nd.array([10])
-    eta, lr, wd, epsilon = 1, 1, 0, 1e-8
-    beta1, beta2 = 0.9, 0.999
-    kwargs = {'eta': eta, 'lr': lr, 'wd': wd, 'epsilon': epsilon,
-              'beta1': beta1, 'beta2': beta2}
-
-    # update is skipped for rescale = nan scalar
-    mx.nd.contrib.adamw_update(weight, grad, m, v,
-                               np.nan, out=weight, **kwargs)
-    # weight remains unchanged
-    mx.test_utils.assert_almost_equal(weight_ref.asnumpy(), weight.asnumpy())
-
-    # update is skipped for rescale = 0
-    mx.nd.contrib.adamw_update(weight, grad, m, v,
-                               rescale_grad * 0, out=weight, **kwargs)
-    # weight remains unchanged
-    mx.test_utils.assert_almost_equal(weight_ref.asnumpy(), weight.asnumpy())
-
-    # update is skipped for rescale = nan
-    mx.nd.contrib.adamw_update(weight, grad, m, v,
-                               rescale_grad * np.nan, out=weight, **kwargs)
-    # weight remains unchanged
-    mx.test_utils.assert_almost_equal(weight_ref.asnumpy(), weight.asnumpy())
-
-    # update is skipped for rescale = inf
-    mx.nd.contrib.adamw_update(weight, grad, m, v,
-                               rescale_grad * np.inf, out=weight, **kwargs)
-    # weight remains unchanged
-    mx.test_utils.assert_almost_equal(weight_ref.asnumpy(), weight.asnumpy())
-
-    # multi-precision update is skipped for rescale = nan
-    weight_fp16 = weight.astype('float16')
-    grad_fp16 = grad.astype('float16')
-    weight_fp16_ref = weight_fp16.copy()
-    mx.nd.contrib.mp_adamw_update(weight_fp16, grad_fp16, m, v, weight,
-                                  rescale_grad * np.nan, out=weight_fp16, **kwargs)
-    mx.test_utils.assert_almost_equal(weight_ref.asnumpy(), weight.asnumpy())
-    mx.test_utils.assert_almost_equal(weight_fp16_ref.asnumpy(), weight_fp16.asnumpy())
-
-    # multi-precision update is skipped for rescale = nan scalar
-    mx.nd.contrib.mp_adamw_update(weight_fp16, grad_fp16, m, v, weight,
-                                  np.nan, out=weight_fp16, **kwargs)
-    mx.test_utils.assert_almost_equal(weight_ref.asnumpy(), weight.asnumpy())
-    mx.test_utils.assert_almost_equal(weight_fp16_ref.asnumpy(), weight_fp16.asnumpy())
-
-    # multi-precision update is skipped for rescale = inf
-    mx.nd.contrib.mp_adamw_update(weight_fp16, grad_fp16, m, v, weight,
-                                  rescale_grad * np.inf, out=weight_fp16, **kwargs)
-    mx.test_utils.assert_almost_equal(weight_ref.asnumpy(), weight.asnumpy())
-    mx.test_utils.assert_almost_equal(weight_fp16_ref.asnumpy(), weight_fp16.asnumpy())
-
-    # multi-precision update is skipped for rescale = 0
-    mx.nd.contrib.mp_adamw_update(weight_fp16, grad_fp16, m, v, weight,
-                                  rescale_grad * 0, out=weight_fp16, **kwargs)
-    mx.test_utils.assert_almost_equal(weight_ref.asnumpy(), weight.asnumpy())
-    mx.test_utils.assert_almost_equal(weight_fp16_ref.asnumpy(), weight_fp16.asnumpy())
-
-    # reference normal update
-    grad_rescale = rescale_grad * grad
-    m_ref = beta1*m + (1-beta1)*grad_rescale
-    v_ref = beta2*v + (1-beta2)*(grad_rescale**2)
-    weight_ref = weight - eta * (1 * m_ref / (v_ref.sqrt() + epsilon) + weight * wd)
-    m_test = m.copy()
-    v_test = v.copy()
-    weight_test = weight.copy()
-    # op normal update
-    mx.nd.contrib.adamw_update(weight_test, grad, m_test, v_test,
-                               rescale_grad, out=weight_test, **kwargs)
-    mx.test_utils.assert_almost_equal(weight_ref.asnumpy(), weight_test.asnumpy())
-    mx.test_utils.assert_almost_equal(m_ref.asnumpy(), m_test.asnumpy())
-    mx.test_utils.assert_almost_equal(v_ref.asnumpy(), v_test.asnumpy())
-
-    # reference normal multi-precision update
-    m_fp32 = m.copy()
-    v_fp32 = v.copy()
-    weight_fp32 = weight.copy()
-    grad_rescale = rescale_grad * grad_fp16.astype('float32')
-    m_ref = beta1*m_fp32 + (1-beta1)*grad_rescale
-    v_ref = beta2*v_fp32 + (1-beta2)*(grad_rescale**2)
-    weight_ref = weight - eta * (1 * m_ref / (v_ref.sqrt() + epsilon) + weight * wd)
-    weight_fp16_ref = weight_ref.astype('float16')
-    # op normal multi-precision update
-    mx.nd.contrib.mp_adamw_update(weight_fp16, grad_fp16, m_fp32, v_fp32, weight_fp32,
-                                  rescale_grad, out=weight_fp16, **kwargs)
-    mx.test_utils.assert_almost_equal(m_ref.asnumpy(), m_fp32.asnumpy())
-    mx.test_utils.assert_almost_equal(v_ref.asnumpy(), v_fp32.asnumpy())
-    mx.test_utils.assert_almost_equal(weight_ref.asnumpy(), weight_fp32.asnumpy())
-    mx.test_utils.assert_almost_equal(weight_fp16_ref.asnumpy(), weight_fp16.asnumpy())
 
+@with_seed()
+def test_adamw():
+    def get_refs(m, v, weight, grad_rescale, beta1, beta2, lr, eta, wd, epsilon, clip_grad=-1):
+        if clip_grad >= 0:
+            grad_rescale = mx.nd.clip(grad_rescale, -clip_grad, clip_grad)
+
+        mean_ref = beta1*m + (1-beta1)*grad_rescale
+        v_ref = beta2*v + (1-beta2)*(grad_rescale**2)
+        weight_ref = weight - eta * (lr * mean_ref / (v_ref.sqrt() + epsilon) + weight * wd)
+        return mean_ref, v_ref, weight_ref
+
+    def run_adamw_test(nElem=1, aggregate=False):
+        aggregate = aggregate or nElem > 1
+        rescale_factor = 10
+        eta, lr, wd, epsilon = 1, 1, 0.1, 1e-8
+        beta1, beta2 = 0.9, 0.999
+        clip_gradient = np.random.uniform(rescale_factor, rescale_factor)
+        weight, grad, m, v, etas, lrs, wds, weight_ref = [], [], [], [], [], [], [], []
+        for i in range(nElem):
+            shape = (np.random.randint(3, high=10), np.random.randint(3, high=10))
+            weight.append(mx.nd.random.uniform(shape=shape))
+            grad.append(mx.nd.random.uniform(-1.0, 1.0, shape=shape))
+            m.append(mx.nd.random.uniform(shape=shape))
+            v.append(mx.nd.random.uniform(shape=shape))
+            etas.append(eta - 1 / np.random.uniform(9, 10))
+            lrs.append(lr - 1 / np.random.uniform(9, 10))
+            wds.append(wd - 1 / np.random.uniform(95, 105))
+            weight_ref.append(weight[i].copy())
+
+        if aggregate:
+            kwargs = {'etas': etas, 'lrs': lrs, 'wds': wds}
+        else:
+            kwargs = {'eta': etas[0], 'lr': lrs[0], 'wd': wds[0]}
+
+        kwargs.update([('epsilon', epsilon), ('beta1', beta1), ('beta2', beta2), ('clip_gradient', clip_gradient)])
+
+        # Test 1: Update is skipped for rescale = nan scalar
+        rescale_grad = mx.nd.array([rescale_factor])
+        tested_grad = [rescale_grad * 0, rescale_grad * np.nan, rescale_grad * np.inf]
+        tested_rescaled_grad = [np.nan]
+        tested_rescaled_grad.extend(tested_grad)
+
+        for rescaled_grad in tested_rescaled_grad:
+            if aggregate:
+                mx.nd.contrib.multi_adamw_update(weight, grad, m, v,
+                                                 rescaled_grad, out=weight, **kwargs)
+            else:
+                mx.nd.contrib.adamw_update(weight[0], grad[0], m[0], v[0],
+                                           rescaled_grad, out=weight[0], **kwargs)
+
+            # weights should remain unchanged
+            for j in range(nElem):
+                assert_almost_equal(weight_ref[j].asnumpy(), weight[j].asnumpy())
+
+
+        # Test 2: Same as Test 1 for multi-precision update
+        weight_fp16, grad_fp16, weight_fp16_refs = [], [], []
+        for i in range(nElem):
+            weight_fp16.append(weight[i].astype('float16'))
+            grad_fp16.append(grad[i].astype('float16'))
+            weight_fp16_refs.append(weight_fp16[i].copy())
+
+        for rescaled_grad in tested_grad:
+            if aggregate:
+                mx.nd.contrib.multi_mp_adamw_update(weight_fp16, grad_fp16, m, v, weight,
+                                                    rescaled_grad, out=weight_fp16, **kwargs)
+            else:
+                mx.nd.contrib.mp_adamw_update(weight_fp16[0], grad_fp16[0], m[0], v[0], weight[0],
+                                              rescaled_grad, out=weight_fp16[0], **kwargs)
+
+            # weights should remain unchanged
+            for i in range(nElem):
+                assert_almost_equal(weight_ref[i].asnumpy(), weight[i].asnumpy())
+                assert_almost_equal(weight_fp16_refs[i].asnumpy(), weight_fp16[i].asnumpy())
+
+
+        # Test 3: Reference normal update
+        grad_rescale, weight_test, m_refs, v_refs, weight_refs = [], [], [], [], []
+        for i in range(nElem):
+            grad_rescale.append(rescale_grad * grad[i])
+            m_ref, v_ref, weight_ref = get_refs(m[i], v[i], weight[i], grad_rescale[i], beta1, beta2, lrs[i], etas[i], wds[i], epsilon, clip_gradient)
+            m_refs.append(m_ref)
+            v_refs.append(v_ref)
+            weight_refs.append(weight_ref)
+            weight_test.append(weight[i].copy())
+
+        # op normal update
+        if aggregate:
+            mx.nd.contrib.multi_adamw_update(weight_test, grad, m, v,
+                                             rescale_grad, out=weight_test, **kwargs)
+        else:
+            mx.nd.contrib.adamw_update(weight_test[0], grad[0], m[0], v[0],
+                                       rescale_grad, out=weight_test[0], **kwargs)
+
+        # Compare results
+        atol = 1e-4 if aggregate else 1e-5
+        rtol = 1e-4 if aggregate else None
+        for i in range(nElem):
+            assert_almost_equal(weight_refs[i].asnumpy(), weight_test[i].asnumpy(), rtol=rtol, atol=atol)
+            assert_almost_equal(m_refs[i].asnumpy(), m[i].asnumpy(), rtol=rtol, atol=atol)
+            assert_almost_equal(v_refs[i].asnumpy(), v[i].asnumpy(), atol=atol)
+
+
+        # Test 4: Reference normal multi-precision update
+        grad_rescale, m_refs, v_refs, weight_refs, weight_fp16_refs = [], [], [], [], []
+        for i in range(nElem):
+            grad_rescale.append(rescale_grad * grad_fp16[i].astype('float32'))
+            m_ref, v_ref, weight_ref = get_refs(m[i], v[i], weight[i], grad_rescale[i], beta1, beta2, lrs[i], etas[i], wds[i], epsilon, clip_gradient)
+            m_refs.append(m_ref)
+            v_refs.append(v_ref)
+            weight_refs.append(weight_ref)
+            weight_fp16_refs.append(weight_ref.astype('float16'))
+
+        # op normal multi-precision update
+        if aggregate:
+            mx.nd.contrib.multi_mp_adamw_update(weight_fp16, grad_fp16, m, v, weight,
+                                                rescale_grad, out=weight_fp16, **kwargs)
+        else:
+            mx.nd.contrib.mp_adamw_update(weight_fp16[0], grad_fp16[0], m[0], v[0], weight[0],
+                                          rescale_grad, out=weight_fp16[0], **kwargs)
+
+        # Compare results
+        for i in range(nElem):
+            assert_almost_equal(m_refs[i].asnumpy(), m[i].asnumpy(), rtol=rtol, atol=atol)
+            assert_almost_equal(v_refs[i].asnumpy(), v[i].asnumpy(), atol=atol)
+            assert_almost_equal(weight_refs[i].asnumpy(), weight[i].asnumpy(), rtol=rtol, atol=atol)
+            assert_almost_equal(weight_fp16_refs[i].asnumpy(), weight_fp16[i].asnumpy(), rtol=1e-3, atol=atol)
+
+    # Testing aggregated Adam update for one element
+    run_adamw_test(1, aggregate=True)
+
+    # Testing Adam update, if nElem = 0, OR
+    #         aggregated Adam update, if nElem > 0
+    for nElem in range(6):
+        run_adamw_test(nElem+1)
 
 if __name__ == '__main__':
     import nose

From 98bf030cd4357762d8cf7edd0cbe67d3ef1e9fff Mon Sep 17 00:00:00 2001
From: andreii <andreii@nvidia.com>
Date: Wed, 9 Oct 2019 09:52:14 -0700
Subject: [PATCH 3/9] Fixing problem with getRescaleGrad(...) call in Python2
 and some minor changes requested by Przemek

---
 python/mxnet/ndarray/contrib.py  | 13 ++++++++-----
 src/operator/contrib/adamw-inl.h |  4 ++--
 2 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/python/mxnet/ndarray/contrib.py b/python/mxnet/ndarray/contrib.py
index d0eac246f4bf..d94bb63c02ae 100644
--- a/python/mxnet/ndarray/contrib.py
+++ b/python/mxnet/ndarray/contrib.py
@@ -21,6 +21,7 @@
 from __future__ import absolute_import
 import math
 import numpy as np
+import mxnet as mx
 from ..context import current_context
 from ..random import uniform
 from ..base import _as_list
@@ -547,7 +548,7 @@ def isnan(data):
     """
     return data != data  # pylint: disable=comparison-with-itself
 
-def getRescaleGrad(rescale_grad, ctx):
+def getRescaleGrad(rescale_grad, ctx=mx.cpu()):
     if not isinstance(rescale_grad, ndarray.NDArray):
         return ndarray.full(shape=(1,), val=rescale_grad, ctx=ctx)
     else:
@@ -555,7 +556,7 @@ def getRescaleGrad(rescale_grad, 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)
+    rescale_grad = getRescaleGrad(rescale_grad, ctx=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,
@@ -565,7 +566,7 @@ 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):
-    rescale_grad = getRescaleGrad(rescale_grad, weight.context)
+    rescale_grad = getRescaleGrad(rescale_grad, ctx=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,
@@ -575,10 +576,11 @@ def mp_adamw_update(weight, grad, mean, var, weight32, rescale_grad, lr, eta, be
 
 def multi_adamw_update(weights, grads, mean, var, rescale_grad, lrs, wds, etas,
                        out=None, name=None, size=0, **kwargs):
+    rescale_grad = getRescaleGrad(rescale_grad, ctx=weights[0].context)
     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),
+                                                 rescale_grad,
                                                  out=out,
                                                  num_weights=size,
                                                  lrs=lrs,
@@ -589,11 +591,12 @@ def multi_adamw_update(weights, grads, mean, var, rescale_grad, lrs, wds, etas,
 
 def multi_mp_adamw_update(weights, grads, mean, var, weights32, rescale_grad, lrs, wds, etas,
                           out=None, name=None, size=0, **kwargs):
+    rescale_grad = getRescaleGrad(rescale_grad, ctx=weights[0].context)
     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),
+        rescale_grad,
         out=out,
         num_weights=size,
         lrs=lrs,
diff --git a/src/operator/contrib/adamw-inl.h b/src/operator/contrib/adamw-inl.h
index 3ff7bf4d0d32..6de8c83f7c41 100644
--- a/src/operator/contrib/adamw-inl.h
+++ b/src/operator/contrib/adamw-inl.h
@@ -135,7 +135,7 @@ struct MPAdamWUpdate {
                const std::vector<TBlob> &outputs,
                const float rescale_grad) {
     using namespace mxnet_op;
-    const auto &param = nnvm::get<AdamWParam>(attrs.parsed);
+    const auto& param = nnvm::get<AdamWParam>(attrs.parsed);
     Stream<xpu>* s = ctx.get_stream<xpu>();
     MSHADOW_REAL_TYPE_SWITCH(inputs[0].type_flag_, DType, {
       Tensor<xpu, 2, DType> weight = inputs[0].FlatTo2D<xpu, DType>(s);
@@ -405,7 +405,7 @@ void FillMultiAdamKernelParam(const nnvm::NodeAttrs& attrs,
 
   pParam->count = p.num_weights;
   pParam->max_size = 0;
-  const bool isSame = std::is_same<DType, MPDType>::value;
+  constexpr bool isSame = std::is_same<DType, MPDType>::value;
   for (int i = 0; i < pParam->count; ++i) {
     const auto idx = i * input_stride;
     pParam->sizes[i] = inputs[idx].shape_.Size();

From 960e948fca0f2dce794ef5c0d68c4c04634b9529 Mon Sep 17 00:00:00 2001
From: andreii <andreii@nvidia.com>
Date: Thu, 10 Oct 2019 16:07:31 -0700
Subject: [PATCH 4/9] Fix a problem appearing in Python2

---
 python/mxnet/ndarray/contrib.py | 29 +++++++++++++++--------------
 1 file changed, 15 insertions(+), 14 deletions(-)

diff --git a/python/mxnet/ndarray/contrib.py b/python/mxnet/ndarray/contrib.py
index d94bb63c02ae..567dda1ace0d 100644
--- a/python/mxnet/ndarray/contrib.py
+++ b/python/mxnet/ndarray/contrib.py
@@ -576,11 +576,12 @@ def mp_adamw_update(weight, grad, mean, var, weight32, rescale_grad, lr, eta, be
 
 def multi_adamw_update(weights, grads, mean, var, rescale_grad, lrs, wds, etas,
                        out=None, name=None, size=0, **kwargs):
-    rescale_grad = getRescaleGrad(rescale_grad, ctx=weights[0].context)
     if not size:
         size = len(weights)
-    return ndarray._internal._multi_adamw_update(*_flatten_list(zip(weights, grads, mean, var)),
-                                                 rescale_grad,
+
+    rescale_grad = getRescaleGrad(rescale_grad, ctx=weights[0].context)
+    temp_list = _flatten_list(zip(weights, grads, mean, var)) + [rescale_grad]
+    return ndarray._internal._multi_adamw_update(*temp_list,
                                                  out=out,
                                                  num_weights=size,
                                                  lrs=lrs,
@@ -591,16 +592,16 @@ def multi_adamw_update(weights, grads, mean, var, rescale_grad, lrs, wds, etas,
 
 def multi_mp_adamw_update(weights, grads, mean, var, weights32, rescale_grad, lrs, wds, etas,
                           out=None, name=None, size=0, **kwargs):
-    rescale_grad = getRescaleGrad(rescale_grad, ctx=weights[0].context)
     if not size:
         size = len(weights)
-    return ndarray._internal._multi_mp_adamw_update(
-        *_flatten_list(zip(weights, grads, mean, var, weights32)),
-        rescale_grad,
-        out=out,
-        num_weights=size,
-        lrs=lrs,
-        wds=wds,
-        etas=etas,
-        name=name,
-        **kwargs)
+
+    rescale_grad = getRescaleGrad(rescale_grad, ctx=weights[0].context)
+    temp_list = _flatten_list(zip(weights, grads, mean, var, weights32)) + [rescale_grad]
+    return ndarray._internal._multi_mp_adamw_update(*temp_list,
+                                                    out=out,
+                                                    num_weights=size,
+                                                    lrs=lrs,
+                                                    wds=wds,
+                                                    etas=etas,
+                                                    name=name,
+                                                    **kwargs)

From 4a8f39c4ded147d4f38cfe443d21a7e3224ff646 Mon Sep 17 00:00:00 2001
From: andreii <andreii@nvidia.com>
Date: Fri, 11 Oct 2019 08:26:22 -0700
Subject: [PATCH 5/9] Minor cleanup

---
 python/mxnet/ndarray/contrib.py  | 10 +++++-----
 src/operator/contrib/adamw-inl.h | 11 ++---------
 2 files changed, 7 insertions(+), 14 deletions(-)

diff --git a/python/mxnet/ndarray/contrib.py b/python/mxnet/ndarray/contrib.py
index 567dda1ace0d..d25925da6503 100644
--- a/python/mxnet/ndarray/contrib.py
+++ b/python/mxnet/ndarray/contrib.py
@@ -548,7 +548,7 @@ def isnan(data):
     """
     return data != data  # pylint: disable=comparison-with-itself
 
-def getRescaleGrad(rescale_grad, ctx=mx.cpu()):
+def _getRescaleGrad(rescale_grad, ctx=mx.cpu()):
     if not isinstance(rescale_grad, ndarray.NDArray):
         return ndarray.full(shape=(1,), val=rescale_grad, ctx=ctx)
     else:
@@ -556,7 +556,7 @@ def getRescaleGrad(rescale_grad, ctx=mx.cpu()):
 
 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, ctx=weight.context)
+    rescale_grad = _getRescaleGrad(rescale_grad, ctx=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,
@@ -566,7 +566,7 @@ 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):
-    rescale_grad = getRescaleGrad(rescale_grad, ctx=weight.context)
+    rescale_grad = _getRescaleGrad(rescale_grad, ctx=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,
@@ -579,7 +579,7 @@ def multi_adamw_update(weights, grads, mean, var, rescale_grad, lrs, wds, etas,
     if not size:
         size = len(weights)
 
-    rescale_grad = getRescaleGrad(rescale_grad, ctx=weights[0].context)
+    rescale_grad = _getRescaleGrad(rescale_grad, ctx=weights[0].context)
     temp_list = _flatten_list(zip(weights, grads, mean, var)) + [rescale_grad]
     return ndarray._internal._multi_adamw_update(*temp_list,
                                                  out=out,
@@ -595,7 +595,7 @@ def multi_mp_adamw_update(weights, grads, mean, var, weights32, rescale_grad, lr
     if not size:
         size = len(weights)
 
-    rescale_grad = getRescaleGrad(rescale_grad, ctx=weights[0].context)
+    rescale_grad = _getRescaleGrad(rescale_grad, ctx=weights[0].context)
     temp_list = _flatten_list(zip(weights, grads, mean, var, weights32)) + [rescale_grad]
     return ndarray._internal._multi_mp_adamw_update(*temp_list,
                                                     out=out,
diff --git a/src/operator/contrib/adamw-inl.h b/src/operator/contrib/adamw-inl.h
index 6de8c83f7c41..fd139de3390f 100644
--- a/src/operator/contrib/adamw-inl.h
+++ b/src/operator/contrib/adamw-inl.h
@@ -78,17 +78,12 @@ inline bool MPUpdateInferShape(const nnvm::NodeAttrs& attrs,
                                mxnet::ShapeVector *out_attrs) {
   CHECK_EQ(in_attrs->size(), static_cast<size_t>(total_in)) << " in operator " << attrs.name;
   CHECK_EQ(out_attrs->size(), static_cast<size_t>(n_out)) << " in operator " << attrs.name;
-  // rescale_grad.shape = ()
   SHAPE_ASSIGN_CHECK(*in_attrs, total_in - 1, mxnet::TShape());
   // TODO(@reminisce): change "none" behavior in ElemwiseAttr
   return ElemwiseAttr<mxnet::TShape, shape_is_none, shape_assign, true, shape_string, n_in, n_out>(
       attrs, in_attrs, out_attrs, mxnet::TShape());
 }
 
-// rescale_grad is a reserved argument at position -1. Example:
-// n_in = 2: weight, grad (fp16)
-// n_out = 1: weight (fp16)
-// total_in = 6: weight, grad, mean, var, weight32, rescale_grad (fp32)
 template<int n_in, int n_out, int total_in>
 inline bool MPUpdateInferType(const nnvm::NodeAttrs& attrs,
                               std::vector<int> *in_attrs,
@@ -195,7 +190,6 @@ struct AdamWUpdate {
 ////
 // Multiple gradients in single kernel
 ////
-
 struct MultiAdamWParam : public dmlc::Parameter<MultiAdamWParam> {
   mxnet::Tuple<float> lrs;
   mxnet::Tuple<float> wds;
@@ -203,7 +197,6 @@ struct MultiAdamWParam : public dmlc::Parameter<MultiAdamWParam> {
   float beta1;
   float beta2;
   float epsilon;
-  float rescale_grad;
   float clip_gradient;
   int num_weights;
   DMLC_DECLARE_PARAMETER(MultiAdamWParam) {
@@ -273,7 +266,7 @@ inline bool MP_MultiAdamW_InferShape(const nnvm::NodeAttrs& attrs,
     }
     all_inferred = all_inferred && ElemwiseShape<input_stride, 1>(attrs, &input_vec, &output_vec);
   }
-  // rescale_grad.shape = ()
+
   SHAPE_ASSIGN_CHECK(*in_attrs, param.num_weights*input_stride, mxnet::TShape());
   return all_inferred;
 }
@@ -312,7 +305,7 @@ inline bool MP_MultiAdamW_InferType(const nnvm::NodeAttrs& attrs,
       TYPE_ASSIGN_CHECK(input_types, input_stride * i + input_stride - 1 - j, mshadow::kFloat32);
     }
   }
-  // rescale_grad.type = ()
+
   TYPE_ASSIGN_CHECK(input_types, param.num_weights*input_stride, mshadow::kFloat32);
   return all_inferred;
 }

From f7252b12e4b5646709cd6a5f305d171e8e8774bd Mon Sep 17 00:00:00 2001
From: andreii <andreii@nvidia.com>
Date: Fri, 11 Oct 2019 14:48:47 -0700
Subject: [PATCH 6/9] Changing function name

---
 python/mxnet/ndarray/contrib.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/python/mxnet/ndarray/contrib.py b/python/mxnet/ndarray/contrib.py
index d25925da6503..0297a1dbe575 100644
--- a/python/mxnet/ndarray/contrib.py
+++ b/python/mxnet/ndarray/contrib.py
@@ -548,7 +548,7 @@ def isnan(data):
     """
     return data != data  # pylint: disable=comparison-with-itself
 
-def _getRescaleGrad(rescale_grad, ctx=mx.cpu()):
+def _get_rescale_grad(rescale_grad, ctx=mx.cpu()):
     if not isinstance(rescale_grad, ndarray.NDArray):
         return ndarray.full(shape=(1,), val=rescale_grad, ctx=ctx)
     else:
@@ -556,7 +556,7 @@ def _getRescaleGrad(rescale_grad, ctx=mx.cpu()):
 
 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, ctx=weight.context)
+    rescale_grad = _get_rescale_grad(rescale_grad, ctx=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,
@@ -566,7 +566,7 @@ 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):
-    rescale_grad = _getRescaleGrad(rescale_grad, ctx=weight.context)
+    rescale_grad = _get_rescale_grad(rescale_grad, ctx=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,
@@ -579,7 +579,7 @@ def multi_adamw_update(weights, grads, mean, var, rescale_grad, lrs, wds, etas,
     if not size:
         size = len(weights)
 
-    rescale_grad = _getRescaleGrad(rescale_grad, ctx=weights[0].context)
+    rescale_grad = _get_rescale_grad(rescale_grad, ctx=weights[0].context)
     temp_list = _flatten_list(zip(weights, grads, mean, var)) + [rescale_grad]
     return ndarray._internal._multi_adamw_update(*temp_list,
                                                  out=out,
@@ -595,7 +595,7 @@ def multi_mp_adamw_update(weights, grads, mean, var, weights32, rescale_grad, lr
     if not size:
         size = len(weights)
 
-    rescale_grad = _getRescaleGrad(rescale_grad, ctx=weights[0].context)
+    rescale_grad = _get_rescale_grad(rescale_grad, ctx=weights[0].context)
     temp_list = _flatten_list(zip(weights, grads, mean, var, weights32)) + [rescale_grad]
     return ndarray._internal._multi_mp_adamw_update(*temp_list,
                                                     out=out,

From 76317372a25f4f7bfb729c4913cdf2b3675210d3 Mon Sep 17 00:00:00 2001
From: andreii <andreii@nvidia.com>
Date: Mon, 14 Oct 2019 08:50:34 -0700
Subject: [PATCH 7/9] Trigger CI


From 5e7e6751b8b04a3fee20f033c1e0b5cbf8db5135 Mon Sep 17 00:00:00 2001
From: andreii <andreii@nvidia.com>
Date: Wed, 16 Oct 2019 09:32:43 -0700
Subject: [PATCH 8/9] Eliminating "asnumpy()" conversion

---
 .../python/unittest/test_contrib_optimizer.py | 20 +++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/tests/python/unittest/test_contrib_optimizer.py b/tests/python/unittest/test_contrib_optimizer.py
index d42999d00c28..7cfd0217aa31 100644
--- a/tests/python/unittest/test_contrib_optimizer.py
+++ b/tests/python/unittest/test_contrib_optimizer.py
@@ -149,7 +149,7 @@ def run_adamw_test(nElem=1, aggregate=False):
 
             # weights should remain unchanged
             for j in range(nElem):
-                assert_almost_equal(weight_ref[j].asnumpy(), weight[j].asnumpy())
+                assert_almost_equal(weight_ref[j], weight[j])
 
 
         # Test 2: Same as Test 1 for multi-precision update
@@ -169,8 +169,8 @@ def run_adamw_test(nElem=1, aggregate=False):
 
             # weights should remain unchanged
             for i in range(nElem):
-                assert_almost_equal(weight_ref[i].asnumpy(), weight[i].asnumpy())
-                assert_almost_equal(weight_fp16_refs[i].asnumpy(), weight_fp16[i].asnumpy())
+                assert_almost_equal(weight_ref[i], weight[i])
+                assert_almost_equal(weight_fp16_refs[i], weight_fp16[i])
 
 
         # Test 3: Reference normal update
@@ -195,9 +195,9 @@ def run_adamw_test(nElem=1, aggregate=False):
         atol = 1e-4 if aggregate else 1e-5
         rtol = 1e-4 if aggregate else None
         for i in range(nElem):
-            assert_almost_equal(weight_refs[i].asnumpy(), weight_test[i].asnumpy(), rtol=rtol, atol=atol)
-            assert_almost_equal(m_refs[i].asnumpy(), m[i].asnumpy(), rtol=rtol, atol=atol)
-            assert_almost_equal(v_refs[i].asnumpy(), v[i].asnumpy(), atol=atol)
+            assert_almost_equal(weight_refs[i], weight_test[i], rtol=rtol, atol=atol)
+            assert_almost_equal(m_refs[i], m[i], rtol=rtol, atol=atol)
+            assert_almost_equal(v_refs[i], v[i], atol=atol)
 
 
         # Test 4: Reference normal multi-precision update
@@ -220,10 +220,10 @@ def run_adamw_test(nElem=1, aggregate=False):
 
         # Compare results
         for i in range(nElem):
-            assert_almost_equal(m_refs[i].asnumpy(), m[i].asnumpy(), rtol=rtol, atol=atol)
-            assert_almost_equal(v_refs[i].asnumpy(), v[i].asnumpy(), atol=atol)
-            assert_almost_equal(weight_refs[i].asnumpy(), weight[i].asnumpy(), rtol=rtol, atol=atol)
-            assert_almost_equal(weight_fp16_refs[i].asnumpy(), weight_fp16[i].asnumpy(), rtol=1e-3, atol=atol)
+            assert_almost_equal(m_refs[i], m[i], rtol=rtol, atol=atol)
+            assert_almost_equal(v_refs[i], v[i], atol=atol)
+            assert_almost_equal(weight_refs[i], weight[i], rtol=rtol, atol=atol)
+            assert_almost_equal(weight_fp16_refs[i], weight_fp16[i], rtol=1e-3, atol=atol)
 
     # Testing aggregated Adam update for one element
     run_adamw_test(1, aggregate=True)

From a43c672a4587c527796ec9d82e1a6ad85e729316 Mon Sep 17 00:00:00 2001
From: andreii <andreii@nvidia.com>
Date: Wed, 16 Oct 2019 16:12:42 -0700
Subject: [PATCH 9/9] Trigger CI