text: # Copyright 2020 Huawei Technologies Co., Ltd
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  # http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  """operator dsl function: sgd"""
  
  import akg
  from akg import tvm, topi
  from akg.utils import kernel_exec as utils
  from akg.utils import validation_check as vc_util
  from akg.utils.dsl_create import TensorUtils
  
  
  NUM_ZERO = 0.0
  
  def sgd_compute(parameters, gradient, learning_rate, accum, momentum, stat,
                  dampening=0.0, weight_decay=0.0, nesterov=False):
      """sgd compute implementation"""
      dtype = parameters.dtype
      if dtype == "float16":
          parameters = topi.cast(parameters, "float32")
          accum = topi.cast(accum, "float32")
          learning_rate = topi.cast(learning_rate, "float32")
          gradient = topi.cast(gradient, "float32")
          momentum = topi.cast(momentum, "float32")
          stat = topi.cast(stat, "float32")
  
      # if weight_decay != 0.0, need compute grad_delta to update gradient
      if weight_decay != 0.0:
          parameters = topi.multiply(parameters, tvm.const(1.0, 'float32'))
          grad_delta = topi.multiply(parameters, weight_decay)
          gradient = topi.add(gradient, grad_delta)
  
      stat_mid = topi.multiply(stat, tvm.const(-1, "float32"))
      stat_act = topi.add(stat_mid, tvm.const(1, "float32"))
  
      dampening_t = topi.multiply(stat_act, dampening)
  
      # update accum
      accum_delta = tvm.compute(accum.shape, lambda *indice: accum(*indice) * momentum[0])
  
      gradient_damp = topi.multiply(gradient, dampening_t)
      accum_t = topi.add(accum_delta, gradient)
      if dampening != 0.0:
          accum_t = topi.subtract(accum_t, gradient_damp)
  
      # update parameters
      if nesterov:
          parameters_delta = tvm.compute(gradient.shape, lambda *indice: gradient(*indice) * learning_rate[0])
          parameters_delta_2 = tvm.compute(accum_t.shape, lambda *indice: accum_t(*indice) * momentum[0])
          parameters_delta_2 = tvm.compute(parameters_delta_2.shape,
                                           lambda *indice: parameters_delta_2(*indice) * learning_rate[0])
          parameters_delta = topi.add(parameters_delta, parameters_delta_2)
          parameters_t = topi.subtract(parameters, parameters_delta)
      else:
          parameters_delta = tvm.compute(accum_t.shape, lambda *indice: accum_t(*indice) * learning_rate[0])
          parameters_t = topi.subtract(parameters, parameters_delta)
  
      # update stat
      stat_t = topi.multiply(stat_act, tvm.const(NUM_ZERO, 'float32'))
  
  
      if dtype == "float16":
          parameters_t = topi.cast(parameters_t, "float16")
          accum_t = topi.cast(accum_t, "float16")
          stat_t = topi.cast(stat_t, "float16")
  
      return parameters_t, accum_t, stat_t
  
  
  @vc_util.check_input_type(akg.tvm.tensor.Tensor, akg.tvm.tensor.Tensor, akg.tvm.tensor.Tensor, akg.tvm.tensor.Tensor,
                            akg.tvm.tensor.Tensor, akg.tvm.tensor.Tensor, (float, type(None)), (float, type(None)),
                            (bool, type(None)))
  def sgd(parameters, gradient, accum, stat, learning_rate, momentum, dampening=0.0, weight_decay=0.0, nesterov=False):
      """
      Update parameters, accum and stat according to the SGD algorithm.
  
      accum = accum * momentum + grad
      if use_nesterov is True:
          parameters -= grad * lr + accum * momentum * lr
      else:
          parameters -= accum * lr
  
      Args:
          parameters (tvm.tensor.Tensor): parameters tensor of float32, float16, to be updated.
          gradient (tvm.tensor.Tensor): gradient tensor of float32, float16.
          accum (tvm.tensor.Tensor): accum tensor of float32, float16, to be updated.
          stat (tvm.tensor.Tensor): stat tensor of float32, float16, to be updated.
          momentum (tvm.tensor.Tensor): momentum tensor of float32, float16, shape must be equal to (1,).
          learning_rate (tvm.tensor.Tensor): learning_rate tensor of float32, float16, shape must be equal to (1,).
          dampening (float): Default value is 0.0.
          weight_decay (float): Default value is 0.0.
          nesterov (bool): Default is False.
  
      Return:
          accum_t (tvm.tensor.Tensor): updated accum with same type and shape as accum.
          stat_t (tvm.tensor.Tensor): updated stat with same type and shape as stat.
          parameters_t (tvm.tensor.Tensor): updated parameters with same type and shape as parameters.
  
      """
      if nesterov and dampening != 0:
          raise ValueError("Nesterov requires zero dampening!")
      if weight_decay < 0:
          raise ValueError("weight_decay must > 0.")
  
      # shape check
      vc_util.elemwise_shape_check(parameters.shape, gradient.shape)
      vc_util.elemwise_shape_check(parameters.shape, accum.shape)
      vc_util.elemwise_shape_check(parameters.shape, stat.shape)
  
      # dtype check
      vc_util.ops_dtype_check([parameters.dtype, gradient.dtype, accum.dtype, stat.dtype],
                              vc_util.DtypeForDavinci.ALL_FLOAT)
  
      parameters_t, accum_t, stat_t = sgd_compute(parameters, gradient, learning_rate, accum, momentum, stat, dampening,
                                                  weight_decay, nesterov)
      parameters_t, binds_info = TensorUtils.inplace_set(parameters, parameters_t, "parameters_buf")
      accum_t, binds_info2 = TensorUtils.inplace_set(accum, accum_t, "accum_buf")
      stat_t, binds_info3 = TensorUtils.inplace_set(stat, stat_t, "stat_buf")
      binds_info.update(binds_info2)
      binds_info.update(binds_info3)
      attrs = {utils.BINDS: binds_info}
  
  
      return parameters_t, accum_t, stat_t, attrs