# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.
'''This module includes a set of optimizers for updating model parameters.
It replaces the old optimizers from optimizer.py'''

from singa import tensor
from singa.tensor import Tensor
from singa import autograd
from . import singa_wrap as singa

from deprecated import deprecated


class DecayScheduler:
    # to be used for decaying learning rate or regularization coefficient or momentum, etc.
    def __init__(self, init_value):
        self.init_value = init_value

    def __call__(self, step):
        assert isinstance(step, Tensor)
        return self.call(step)

    def call(self, step) -> Tensor:
        # step is a Tensor with a single scalar value
        # return the current value as a Tensor
        raise NotImplementedError


class Constant(DecayScheduler):

    def call(self, step: Tensor) -> Tensor:
        # TODO should be an in-place operator
        ret = Tensor((1,), step.device)
        ret.set_value(self.init_value)
        return ret


class ExponentialDecay(DecayScheduler):

    def __init__(self, init_value, decay_steps, decay_rate, staircase=False):
        super(ExponentialDecay, self).__init__(init_value)

        self.decay_steps = decay_steps
        self.decay_rate = decay_rate
        self.staircase = staircase

    def call(self, step):
        if self.staircase:
            s = step // self.decay_steps
        else:
            s = step / self.decay_steps
        ret = Tensor((1,), s.device)
        ret.set_value(self.decay_rate)
        return self.init_value * tensor.pow(ret, s)


class Optimizer(object):
    """Base optimizer.

    Args:
        config (Dict): specify the default values of configurable variables.
    """

    def __init__(self, lr):
        # init lr(could be a constant scalar or a learning rate scheduler)
        if type(lr) == float or type(lr) == int:
            self.lr = Constant(lr)
        elif isinstance(lr, DecayScheduler):
            self.lr = lr
        else:
            raise TypeError("Wrong learning rate type")

        # init step counter
        # TODO change type to int32
        self.step_counter = Tensor((1,), dtype=tensor.float32)
        self.step_counter.set_value(0)
        self.lr_value = self.lr(self.step_counter)

    def get_states(self):
        # skip DecayScheduler as it does not have persistent states
        return {'step_counter': tensor.to_numpy(self.step_counter)[0]}

    def set_states(self, states):
        self.step_counter = Tensor((1,))
        self.step_counter.set_value(states['step_counter'])
        self.lr_value = self.lr(self.step_counter)

    def __call__(self, loss):
        self.call(loss)
        self.step()

    def call(self, loss):
        for p, g in autograd.backward(loss):
            if p.name is None:
                p.name = id(p)
            self.apply(p.name, p, g)

    def step(self):
        """To increment the step counter and update the lr"""
        self.step_counter.data += 1
        lr_value = self.lr(self.step_counter)
        self.lr_value.copy_from(lr_value)

    def apply(self, param_name, param_value, param_grad):
        """Performs a single optimization step.

        Args:
                param_name(String): the name of the param
                param_value(Tensor): param values to be update in-place
                grad(Tensor): param gradients; the values may be updated
                        in this function; cannot use it anymore
        """
        raise NotImplementedError

    @deprecated(
        reason=
        "Update is deprecated, use apply() to do update, refer to apply for more details."
    )
    def update(self, param, grad):
        """Update the param values with given gradients.

        Args:
            param(Tensor): param values to be updated in-place
            grad(Tensor): param gradients; the values may be updated
                    in this function; do not use it anymore
        """
        if param.name is None:
            param.name = id(param)
        self.apply(param.name, param, grad)

    def device_check(self, *inputs):
        flag = inputs[0].device.graph_enabled()
        inputs[0].device.EnableGraph(False)
        x_device = inputs[0].device
        x_dev_id = x_device.id()
        for var in inputs:
            if var.device.id() != x_dev_id:
                var.to_device(x_device)
        inputs[0].device.EnableGraph(flag)

    @deprecated(
        reason=
        "backward_and_update is deprecated, use __call__() to do update, refer to __call__ for more details."
    )
    def backward_and_update(self, loss):
        """Performs backward propagation from the loss and parameter update.

        From the loss, it performs backward propagation to get the gradients
        and do the parameter update.

        Args:
                loss(Tensor): loss is the objective function of the deep learning model
                optimization, e.g. for classification problem it can be the output of the
                softmax_cross_entropy function.
        """
        self.__call__(loss)


class SGD(Optimizer):
    """Implements stochastic gradient descent (optionally with momentum).

    Nesterov momentum is based on the formula from `On the importance of initialization and momentum in deep learning`__.

    Args:
        lr(float): learning rate
        momentum(float, optional): momentum factor(default: 0)
        weight_decay(float, optional): weight decay(L2 penalty)(default: 0)
        dampening(float, optional): dampening for momentum(default: 0)
        nesterov(bool, optional): enables Nesterov momentum(default: False)

    Typical usage example:
        >> > from singa import opt
        >> > optimizer = opt.SGD(lr=0.1, momentum=0.9)
        >> > optimizer.update()

    __ http: // www.cs.toronto.edu / %7Ehinton / absps / momentum.pdf

    .. note::
        The implementation of SGD with Momentum / Nesterov subtly differs from
        Sutskever et. al. and implementations in some other frameworks.

        Considering the specific case of Momentum, the update can be written as

        .. math::
                  v = \rho * v + g \\
                  p = p - lr * v

        where p, g, v and: math: `\rho` denote the parameters, gradient,
        velocity, and momentum respectively.

        This is in contrast to Sutskever et. al. and
        other frameworks which employ an update of the form

        .. math::
             v = \rho * v + lr * g \\
             p = p - v

        The Nesterov version is analogously modified.
    """

    def __init__(self,
                 lr=0.1,
                 momentum=0,
                 dampening=0,
                 weight_decay=0,
                 nesterov=False):
        super(SGD, self).__init__(lr)

        # init momentum
        if type(momentum) == float or type(momentum) == int:
            if momentum < 0.0:
                raise ValueError("Invalid momentum value: {}".format(momentum))
            self.momentum = Constant(momentum)
        elif isinstance(momentum, DecayScheduler):
            self.momentum = momentum
            momentum = momentum.init_value
        else:
            raise TypeError("Wrong momentum type")
        self.mom_value = self.momentum(self.step_counter)

        # init dampening
        if type(dampening) == float or type(dampening) == int:
            self.dampening = Constant(dampening)
        elif isinstance(dampening, DecayScheduler):
            self.dampening = dampening
            dampening = dampening.init_value
        else:
            raise TypeError("Wrong dampening type")
        self.dam_value = self.dampening(self.step_counter)

        # init weight_decay
        if type(weight_decay) == float or type(weight_decay) == int:
            if weight_decay < 0.0:
                raise ValueError(
                    "Invalid weight_decay value: {}".format(weight_decay))
            self.weight_decay = Constant(weight_decay)
        elif isinstance(weight_decay, DecayScheduler):
            self.weight_decay = weight_decay
        else:
            raise TypeError("Wrong weight_decay type")
        self.decay_value = self.weight_decay(self.step_counter)

        # init other params
        self.nesterov = nesterov
        self.moments = dict()

        # check value
        if nesterov and (momentum <= 0 or dampening != 0):
            raise ValueError(
                "Nesterov momentum requires a momentum and zero dampening")

    def apply(self, param_name, param_value, param_grad):
        """Performs a single optimization step.

        Args:
                param_name(String): the name of the param
                param_value(Tensor): param values to be update in-place
                grad(Tensor): param gradients; the values may be updated
                        in this function; cannot use it anymore
        """
        assert param_value.shape == param_grad.shape, ("shape mismatch",
                                                       param_value.shape,
                                                       param_grad.shape)
        self.device_check(param_value, self.step_counter, self.lr_value,
                          self.mom_value, self.dam_value, self.decay_value)

        # TODO add branch operator
        # if self.decay_value != 0:
        if self.weight_decay.init_value != 0:
            singa.Axpy(self.decay_value.data, param_value.data, param_grad.data)

        if self.momentum.init_value != 0:
            if param_name not in self.moments:
                flag = param_value.device.graph_enabled()
                param_value.device.EnableGraph(False)
                self.moments[param_name] = tensor.zeros_like(param_value)
                param_value.device.EnableGraph(flag)

            buf = self.moments[param_name]
            buf *= self.mom_value
            alpha = 1.0 - self.dam_value
            singa.Axpy(alpha.data, param_grad.data, buf.data)

            if self.nesterov:
                singa.Axpy(self.mom_value.data, buf.data, param_grad.data)
            else:
                param_grad = buf

        minus_lr = 0.0 - self.lr_value
        singa.Axpy(minus_lr.data, param_grad.data, param_value.data)

    def step(self):
        # increment step counter, lr and moment
        super().step()
        mom_value = self.momentum(self.step_counter)
        dam_value = self.dampening(self.step_counter)
        decay_value = self.weight_decay(self.step_counter)
        self.mom_value.copy_from(mom_value)
        self.dam_value.copy_from(dam_value)
        self.decay_value.copy_from(decay_value)

    def get_states(self):
        states = super().get_states()
        if self.mom_value > 0:
            states[
                'moments'] = self.moments  # a dict for 1st order moments tensors
        return states

    def set_states(self, states):
        super().set_states(states)
        if 'moments' in states:
            self.moments = states['moments']
            self.mom_value = self.momentum(self.step_counter)


class RMSProp(Optimizer):
    '''RMSProp optimizer.

    See the base Optimizer for all constructor args.

    Args:
        rho (float): float within [0, 1]
        epsilon (float): small value for preventing numeric error
    '''

    def __init__(self, lr=0.1, rho=0.9, epsilon=1e-8, weight_decay=0):
        super(RMSProp, self).__init__(lr)

        # init weight_decay
        if type(weight_decay) == float or type(weight_decay) == int:
            if weight_decay < 0.0:
                raise ValueError(
                    "Invalid weight_decay value: {}".format(weight_decay))
            self.weight_decay = Constant(weight_decay)
        elif isinstance(weight_decay, DecayScheduler):
            self.weight_decay = weight_decay
        else:
            raise TypeError("Wrong weight_decay type")
        self.decay_value = self.weight_decay(self.step_counter)

        # init rho
        if type(rho) == float or type(rho) == int:
            self.rho = Constant(rho)
        elif isinstance(rho, DecayScheduler):
            self.rho = rho
        else:
            raise TypeError("Wrong rho type")
        self.rho_value = self.rho(self.step_counter)

        # init epsilon
        if type(epsilon) == float or type(epsilon) == int:
            self.epsilon = Constant(epsilon)
        elif isinstance(rho, DecayScheduler):
            self.epsilon = epsilon
        else:
            raise TypeError("Wrong epsilon type")
        self.epsilon_value = self.epsilon(self.step_counter)

        # init running average
        self.running_average = dict()

    def apply(self, param_name, param_value, param_grad):
        """Performs a single optimization step.

        Args:
                param_name(String): the name of the param
                param_value(Tensor): param values to be update in-place
                grad(Tensor): param gradients; the values may be updated
                        in this function; cannot use it anymore
        """
        assert param_value.shape == param_grad.shape, ("shape mismatch",
                                                       param_value.shape,
                                                       param_grad.shape)
        self.device_check(param_value, self.step_counter, self.lr_value,
                          self.rho_value, self.epsilon_value, self.decay_value)

        # if self.decay_value != 0:
        if self.weight_decay.init_value != 0:
            singa.Axpy(self.decay_value.data, param_value.data, param_grad.data)

        if param_name not in self.running_average:
            flag = param_value.device.graph_enabled()
            param_value.device.EnableGraph(False)
            self.running_average[param_name] = tensor.zeros_like(param_value)
            param_value.device.EnableGraph(flag)

        # running_average = running_average * rho + param_grad * param_grad * (1 - rho)
        # param_value = param_value - lr * param_grad / sqrt(running_average + epsilon)

        self.running_average[param_name] *= self.rho_value

        tmp1 = singa.Square(param_grad.data)
        tmp2 = 1.0 - self.rho_value
        singa.Axpy(tmp2.data, tmp1, self.running_average[param_name].data)

        minus_lr = 0.0 - self.lr_value
        tmp3 = self.running_average[param_name] + self.epsilon_value
        tmp3 = singa.Sqrt(tmp3.data)
        tmp3 = singa.__div__(param_grad.data, tmp3)

        singa.Axpy(minus_lr.data, tmp3, param_value.data)

    def step(self):
        # increment step counter, lr and moment
        super().step()
        decay_value = self.weight_decay(self.step_counter)
        rho_value = self.rho(self.step_counter)
        epsilon_value = self.epsilon(self.step_counter)
        self.decay_value.copy_from(decay_value)
        self.rho_value.copy_from(rho_value)
        self.epsilon_value.copy_from(epsilon_value)

    def get_states(self):
        states = super().get_states()
        states['running_average'] = self.running_average
        return states

    def set_states(self, states):
        super().set_states(states)
        if 'running_average' in states:
            self.running_average = states['running_average']


class AdaGrad(Optimizer):
    '''AdaGrad optimizer.

    See the base Optimizer for all constructor args.

    Args:
        epsilon (float): small number for preventing numeric error.
    '''

    def __init__(self, lr=0.1, epsilon=1e-8, weight_decay=0):
        super(AdaGrad, self).__init__(lr)

        # init weight_decay
        if type(weight_decay) == float or type(weight_decay) == int:
            if weight_decay < 0.0:
                raise ValueError(
                    "Invalid weight_decay value: {}".format(weight_decay))
            self.weight_decay = Constant(weight_decay)
        elif isinstance(weight_decay, DecayScheduler):
            self.weight_decay = weight_decay
        else:
            raise TypeError("Wrong weight_decay type")
        self.decay_value = self.weight_decay(self.step_counter)

        # init epsilon
        if type(epsilon) == float or type(epsilon) == int:
            self.epsilon = Constant(epsilon)
        elif isinstance(epsilon, DecayScheduler):
            self.epsilon = epsilon
        else:
            raise TypeError("Wrong epsilon type")
        self.epsilon_value = self.epsilon(self.step_counter)

        # init history
        self.history = dict()

    def apply(self, param_name, param_value, param_grad):
        """Performs a single optimization step.

        Args:
                param_name(String): the name of the param
                param_value(Tensor): param values to be update in-place
                grad(Tensor): param gradients; the values may be updated
                        in this function; cannot use it anymore
        """
        assert param_value.shape == param_grad.shape, ("shape mismatch",
                                                       param_value.shape,
                                                       param_grad.shape)
        self.device_check(param_value, self.step_counter, self.lr_value,
                          self.epsilon_value, self.decay_value)

        # if self.decay_value != 0:
        if self.weight_decay.init_value != 0:
            singa.Axpy(self.decay_value.data, param_value.data, param_grad.data)

        if param_name not in self.history:
            flag = param_value.device.graph_enabled()
            param_value.device.EnableGraph(False)
            self.history[param_name] = tensor.zeros_like(param_value)
            param_value.device.EnableGraph(flag)

        # history = history + param_grad * param_grad
        # param_value = param_value - lr * param_grad / sqrt(history + epsilon)

        tmp = self.history[param_name].data
        tmp += singa.Square(param_grad.data)

        minus_lr = 0.0 - self.lr_value
        tmp = self.history[param_name] + self.epsilon_value
        tmp = singa.Sqrt(tmp.data)
        tmp = singa.__div__(param_grad.data, tmp)
        singa.Axpy(minus_lr.data, tmp, param_value.data)

    def step(self):
        # increment step counter, lr and moment
        super().step()
        decay_value = self.weight_decay(self.step_counter)
        epsilon_value = self.epsilon(self.step_counter)
        self.decay_value.copy_from(decay_value)
        self.epsilon_value.copy_from(epsilon_value)

    def get_states(self):
        states = super().get_states()
        states['history'] = self.history  # a dict for 1st order moments tensors
        return states

    def set_states(self, states):
        super().set_states(states)
        if 'history' in states:
            self.history = states['history']


class Adam(Optimizer):
    '''Adam optimizer.

    See the base Optimizer for all constructor args.

    Args:
        beta_1(float): coefficient of momentum
        beta_2(float): coefficient of aggregated squared gradient
        epsilon (float): small value for preventing numeric error
    '''

    def __init__(self,
                 lr=0.1,
                 beta_1=0.9,
                 beta_2=0.999,
                 epsilon=1e-8,
                 weight_decay=0):
        super(Adam, self).__init__(lr)

        # init weight_decay
        if type(weight_decay) == float or type(weight_decay) == int:
            if weight_decay < 0.0:
                raise ValueError(
                    "Invalid weight_decay value: {}".format(weight_decay))
            self.weight_decay = Constant(weight_decay)
        elif isinstance(weight_decay, DecayScheduler):
            self.weight_decay = weight_decay
        else:
            raise TypeError("Wrong weight_decay type")
        self.decay_value = self.weight_decay(self.step_counter)

        # init beta_1
        if type(beta_1) == float or type(beta_1) == int:
            self.beta_1 = Constant(beta_1)
        elif isinstance(beta_1, DecayScheduler):
            self.beta_1 = beta_1
        else:
            raise TypeError("Wrong beta_1 type")
        self.beta_1_value = self.beta_1(self.step_counter)

        # init beta_2
        if type(beta_2) == float or type(beta_2) == int:
            self.beta_2 = Constant(beta_2)
        elif isinstance(beta_2, DecayScheduler):
            self.beta_2 = beta_2
        else:
            raise TypeError("Wrong beta_2 type")
        self.beta_2_value = self.beta_2(self.step_counter)

        # init epsilon
        if type(epsilon) == float or type(epsilon) == int:
            self.epsilon = Constant(epsilon)
        elif isinstance(epsilon, DecayScheduler):
            self.epsilon = epsilon
        else:
            raise TypeError("Wrong epsilon type")
        self.epsilon_value = self.epsilon(self.step_counter)

        # init m and v
        self.m = dict()
        self.v = dict()

    def apply(self, param_name, param_value, param_grad):
        """Performs a single optimization step.

        Args:
                param_name(String): the name of the param
                param_value(Tensor): param values to be update in-place
                grad(Tensor): param gradients; the values may be updated
                        in this function; cannot use it anymore
        """
        assert param_value.shape == param_grad.shape, ("shape mismatch",
                                                       param_value.shape,
                                                       param_grad.shape)
        self.device_check(param_value, self.step_counter, self.lr_value,
                          self.beta_1_value, self.beta_2_value,
                          self.epsilon_value, self.decay_value)

        # if self.decay_value != 0:
        if self.weight_decay.init_value != 0:
            singa.Axpy(self.decay_value.data, param_value.data, param_grad.data)

        if param_name not in self.m:
            flag = param_value.device.graph_enabled()
            param_value.device.EnableGraph(False)
            self.m[param_name] = tensor.zeros_like(param_value)
            self.v[param_name] = tensor.zeros_like(param_value)
            param_value.device.EnableGraph(flag)

        # overall steps
        # m := beta_1 * m + (1 - beta_1) * grad
        # v := beta_2 * v + (1 - beta_2) * grad * grad
        # m_norm = m / (1 - beta_1 ^ step)
        # v_norm = v / (1 - beta_2 ^ step)
        # param := param - (lr * m_norm) / ( sqrt(v_norm) + epsilon) )

        step = self.step_counter + 1.0

        # m := beta_1 * m + (1 - beta_1) * grad
        tmp = 1.0 - self.beta_1_value
        self.m[param_name] *= self.beta_1_value
        singa.Axpy(tmp.data, param_grad.data, self.m[param_name].data)

        # v := beta_2 * v + (1 - beta_2) * grad * grad
        tmp = 1.0 - self.beta_2_value
        self.v[param_name] *= self.beta_2_value
        singa.Axpy(tmp.data, singa.Square(param_grad.data),
                   self.v[param_name].data)

        # m_norm = m / (1 - beta_1 ^ step)
        tmp = tensor.pow(self.beta_1_value, step)
        tmp = 1.0 - tmp
        m_norm = self.m[param_name] / tmp

        # v_norm = v / (1 - beta_2 ^ step)
        tmp = tensor.pow(self.beta_2_value, step)
        tmp = 1.0 - tmp
        v_norm = self.v[param_name] / tmp

        # param := param - (lr * m_norm) / ( sqrt(v_norm) + epsilon) )
        a = tensor.sqrt(v_norm) + self.epsilon_value
        tmp = m_norm / a

        minus_lr = 0.0 - self.lr_value
        singa.Axpy(minus_lr.data, tmp.data, param_value.data)

    def step(self):
        # increment step counter, lr and moment
        super().step()
        decay_value = self.weight_decay(self.step_counter)
        beta_1_value = self.beta_1(self.step_counter)
        beta_2_value = self.beta_2(self.step_counter)
        self.decay_value.copy_from(decay_value)
        self.beta_1_value.copy_from(beta_1_value)
        self.beta_2_value.copy_from(beta_2_value)

    def get_states(self):
        states = super().get_states()
        states['m'] = self.m  # a dict for 1st order moments tensors
        states['v'] = self.v  # a dict for 2nd order moments tensors
        return states

    def set_states(self, states):
        super().set_states(states)
        if 'm' in states:
            self.m = states['m']
        if 'v' in states:
            self.v = states['v']


class DistOpt(object):
    """The class is designed to wrap an optimizer to do distributed training.

    This class is used to wrap an optimizer object to perform distributed training based
    on multiprocessing. Each process has an individual rank, which gives information of
    which GPU the individual process is using. The training data is partitioned, so that
    each process can evaluate the sub-gradient based on the partitioned training data.
    Once the sub-graident is calculated on each processes, the overall stochastic gradient
    is obtained by all-reducing the sub-gradients evaluated by all processes. The all-reduce
    operation is supported by the NVidia Collective Communication Library (NCCL).

    Args:
        opt(Optimizer): The optimizer to be wrapped.
        nccl_id(NcclIdHolder): an nccl id holder object for a unique communication id
        local_rank(int): local rank of a process on the current node
        world_size(int): total number of processes
        buffSize(int): the buffSize in terms of number of elements used in nccl communicator

    Attributes:
        world_size(int): total number of processes
        local_rank(int): local rank of a process on the current node
        global_rank(int): global rank of a process

    Typical usage example:
        >> > from singa import opt
        >> > optimizer = opt.SGD(lr=0.1, momentum=0.9)
        >> > optimizer = opt.DistOpt(sgd)

    """

    def __init__(self,
                 opt=SGD(),
                 nccl_id=None,
                 local_rank=None,
                 world_size=None,
                 buffSize=4194304):
        self.opt = opt
        if nccl_id is None:
            # constructure for application using MPI
            self.communicator = singa.Communicator(buffSize)
        else:
            # constructor for application using python multi-process module
            self.communicator = singa.Communicator(local_rank, world_size,
                                                   nccl_id, buffSize)

        self.world_size = self.communicator.world_size
        self.local_rank = self.communicator.local_rank
        self.global_rank = self.communicator.global_rank

    def __call__(self, loss):
        self.backward_and_update(loss)

    def update(self, param, grad):
        """Performs a single optimization step.

        Args:
                param(Tensor): param values to be update
                grad(Tensor): param gradients
        """
        grad /= self.world_size
        self.opt.update(param, grad)

    def all_reduce(self, tensor):
        """Performs all reduce of a tensor for distributed training.

        Args:
                tensor(Tensor): a tensor to be all-reduced
        """
        self.communicator.synch(tensor)

    def fused_all_reduce(self, tensor, send=True):
        """Performs all reduce of the tensors after fusing them in a buffer.

        Args:
                tensor(List of Tensors): a list of tensors to be all-reduced
                send(bool): When send is False, the tensor won't be send to the
                target device immediately, it will be copied to the buffer first
        """
        tensor = singa.VecTensor(tensor)
        self.communicator.fusedSynch(tensor, send)

    def all_reduce_half(self, tensor):
        """Performs all reduce of a tensor after converting to FP16.

        Args:
                tensor(Tensor): a tensor to be all-reduced
        """
        self.communicator.synchHalf(tensor)

    def fused_all_reduce_half(self, tensor, send=True):
        """Performs all reduce of the tensors after fusing and converting them to FP16.

        Args:
                tensor(List of Tensors): a list of tensors to be all-reduced
                send(bool): When send is False, the tensor won't be send to the
                target device immediately, it will be copied to the buffer first
        """
        tensor = singa.VecTensor(tensor)
        self.communicator.fusedSynchHalf(tensor, send)

    def sparsification(self, tensor, accumulation, spars, topK):
        """Performs all reduce of a tensor after sparsification.

        Args:
                tensor(Tensor): a tensor to be all-reduced
                accumulation(Tensor): local gradient accumulation
                spars(float): a parameter to control sparsity as defined below
                topK(bool): When topK is False, it sparsifies the gradient with absolute
                value >= sparsWhen topK is True, it sparsifies a fraction of total gradient
                number equals to spars,  E.g. when spars = 0.01, it sparsifies 1 % of the
                total gradient elements
        """
        if accumulation is None:
            self.communicator.sparsification(tensor, spars, topK)
        else:
            self.communicator.sparsification(tensor, accumulation, spars, topK)

    def fused_sparsification(self, tensor, accumulation, spars, topK):
        """Performs all reduce of the tensors after fusing and sparsification.

        Args:
                tensor(List of Tensors): a list of tensors to be all-reduced
                accumulation(Tensor): local gradient accumulation
                spars(float): a parameter to control sparsity as defined below
                topK(bool): When topK is False, it sparsifies the gradient with absolute
                value >= sparsWhen topK is True, it sparsifies a fraction of total gradient
                number equals to spars,  E.g. when spars = 0.01, it sparsifies 1 % of the
                total gradient elements
        """
        tensor = singa.VecTensor(tensor)
        if accumulation is None:
            self.communicator.fusedSparsification(tensor, spars, topK)
        else:
            self.communicator.fusedSparsification(tensor, accumulation, spars,
                                                  topK)

    def wait(self):
        """Wait for the cuda streams used by the communicator to finish their operations."""
        self.communicator.wait()

    def backward_and_update(self, loss, threshold=2097152):
        """Performs backward propagation from the loss and parameter update.

        From the loss, it performs backward propagation to get the gradients and do the parameter
        update. For gradient communication, it fuses all the tensor smaller than the threshold
        value to reduce network latency.

        Args:
                loss(Tensor): loss is the objective function of the deep learning model
                optimization, e.g. for classification problem it can be the output of the
                softmax_cross_entropy function.
                threshold(int): threshold is a parameter to control performance in fusing
                the tensors. For the tensors of sizes smaller than threshold, they are to
                be accumulated and fused before the all reduce operation. For the tensors
                of its size larger than the threshold value, they are to be reduced directly
                without fusion.
        """
        plist = []
        acc = 0
        glist = []
        for p, g in autograd.backward(loss):
            if g.size() > threshold:
                # larger than threshold -> reduced directly
                self.all_reduce(g.data)
            else:
                # smaller than threshold -> accumulate
                glist.append(g.data)
                self.fused_all_reduce([g.data], send=False)
                acc += g.size()
                if (acc > threshold):
                    self.fused_all_reduce(glist)
                    acc = 0
                    glist = []
            plist.append((p, g))
        if glist:
            self.fused_all_reduce(glist)
        self.wait()
        for p, g in plist:
            self.update(p, g)
        self.opt.step()

    def backward_and_update_half(self,
                                 loss,
                                 threshold=2097152,
                                 clipping=False,
                                 clip_Value=100):
        """Performs backward propagation and parameter update, with FP16 precision communication.

        THIS IS A EXPERIMENTAL FUNCTION FOR RESEARCH PURPOSE:
        From the loss, it performs backward propagation to get the gradients and do the parameter
        update. For gradient communication, it fuses all the tensor smaller than the threshold value
        to reduce network latency, as well as converting them to FP16 half precision format before
        sending them out. To assist training, this functions provide an option to perform gradient
        clipping.

        Args:
                loss(Tensor): loss is the objective function of the deep learning model
                optimization, e.g. for classification problem it can be the output of the
                softmax_cross_entropy function.
                threshold(int): threshold is a parameter to control performance in fusing
                the tensors. For the tensors of sizes smaller than threshold, they are to
                be accumulated and fused before the all reduce operation. For the tensors
                of its size larger than the threshold value, they are to be reduced directly
                without fusion.
                clipping(bool): a boolean flag to choose whether to clip the gradient value
                clip_value(float): the clip value to be used when clipping is True
        """
        plist = []
        acc = 0
        glist = []
        for p, g in autograd.backward(loss):
            if clipping:
                g = autograd.clip(g, -clip_Value, clip_Value)
            if g.size() > threshold:
                # larger than threshold -> reduced directly
                self.all_reduce_half(g.data)
            else:
                # smaller than threshold -> accumulate
                glist.append(g.data)
                self.fused_all_reduce_half([g.data], send=False)
                acc += g.size()
                if (acc > threshold):
                    self.fused_all_reduce_half(glist)
                    acc = 0
                    glist = []
            plist.append((p, g))
        if glist:
            self.fused_all_reduce_half(glist)
        self.wait()
        for p, g in plist:
            self.update(p, g)
        self.opt.step()

    def backward_and_partial_update(self, loss, threshold=2097152):
        """Performs backward propagation from the loss and parameter update using asychronous training.

        THIS IS A EXPERIMENTAL FUNCTION FOR RESEARCH PURPOSE:
        From the loss, it performs backward propagation to get the gradients and do the parameter
        update. It fuses the tensors smaller than the threshold value to reduce network latency,
        as well as performing asychronous training where one parameter partition is all-reduced
        per iteration. The size of the parameter partition depends on the threshold value.

        Args:
                loss(Tensor): loss is the objective function of the deep learning model
                optimization, e.g. for classification problem it can be the output of the
                softmax_cross_entropy function.
                threshold(int): threshold is a parameter to control performance in fusing
                the tensors. For the tensors of sizes smaller than threshold, they are to
                be accumulated and fused before the all reduce operation. For the tensors
                of its size larger than the threshold value, they are to be reduced directly
                without fusion.

        Attributes:
                self.partial(int): A counter to determine which partition to perform all-reduce.
                This counter resets to zero automatlly after an update cycle of the full parameter
                set.
        """
        if not hasattr(self, "partial"):
            self.partial = 0
        self.partial += 1
        k = 0
        plist = []
        acc = 0
        tenlist = []
        reduced = []
        for p, g in autograd.backward(loss):
            # every parameters update locally
            self.opt.update(p, g)
            # then do the partial parameter sychronization
            if p.size() > threshold:
                # larger than threshold -> reduced directly
                # k is the partition number of the full gradient set
                k += 1
                if (k == self.partial):
                    self.all_reduce(p.data)
                    reduced.append(p)
            else:
                # smaller than threshold -> accumulate
                plist.append(p.data)
                tenlist.append(p)
                acc += p.size()
                if (acc > threshold):
                    k += 1
                    if (k == self.partial):
                        self.fused_all_reduce(plist, send=False)
                        self.fused_all_reduce(plist)
                        reduced = tenlist
                    acc = 0
                    plist = []
                    tenlist = []
        if plist:
            k += 1
            if (k == self.partial):
                self.fused_all_reduce(plist, send=False)
                self.fused_all_reduce(plist)
                reduced = tenlist
        self.wait()
        # the all-reduced parameters needed to be averaged
        for r in reduced:
            r /= self.world_size
        # the counter returns to zero after a cycle of partial update
        if (k == self.partial):
            self.partial = 0
        self.opt.step()

    def backward_and_sparse_update(self,
                                   loss,
                                   threshold=2097152,
                                   spars=0.05,
                                   topK=False,
                                   corr=True):
        """ Performs backward propagation from the loss and parameter update with sparsification.

        THIS IS A EXPERIMENTAL FUNCTION FOR RESEARCH PURPOSE:
        From the loss, it performs backward propagation to get the gradients and do the parameter
        update. It fuses the tensors with size smaller than the threshold value to reduce network
        latency, as well as using sparsification schemes to transfer only the gradient elements which
        are significant.

        Args:
                loss(Tensor): loss is the objective function of the deep learning model
                optimization, e.g. for classification problem it can be the output of the
                softmax_cross_entropy function.
                threshold(int): threshold is a parameter to control performance in fusing
                the tensors. For the tensors of sizes smaller than threshold, they are to
                be accumulated and fused before the all reduce operation. For the tensors
                of its size larger than the threshold value, they are to be reduced directly
                without fusion.
                spars(float): a parameter to control sparsity as defined below
                topK(bool): When topK is False, it sparsifies the gradient with absolute
                value >= sparsWhen topK is True, it sparsifies a fraction of total gradient
                number equals to spars,  E.g. when spars = 0.01, it sparsifies 1 % of the
                total gradient elements
                corr(bool): whether to use the local accumulate gradient for correction

        Attributes:
                self.sparsInit: A counter to determine which partition to perform all-reduce.
                self.gradAccumulation: Local gradient accumulation
        """
        if ((not hasattr(self, "sparsInit")) and corr):
            self.gradAccumulation = []
            self.sparsInit = False
        plist = []
        acc = 0
        k = -1
        glist = []
        for p, g in autograd.backward(loss):
            if g.size() > threshold:
                # larger than threshold -> reduced directly
                k += 1
                if (corr and (not self.sparsInit)):
                    # create a tensor for the gradient accumulation
                    flag = p.device.graph_enabled()
                    p.device.EnableGraph(False)
                    self.gradAccumulation.append(
                        tensor.Tensor((g.size(),), p.device, p.dtype))
                    self.gradAccumulation[k].set_value(0.0)
                    p.device.EnableGraph(flag)
                if corr:
                    self.sparsification(g.data, self.gradAccumulation[k].data,
                                        spars, topK)
                else:
                    self.sparsification(g.data, None, spars, topK)
            else:
                # smaller than threshold -> accumulate
                glist.append(g.data)
                acc += g.size()
                if (acc > threshold):
                    k += 1
                    if (corr and (not self.sparsInit)):
                        # create a tensor for the gradient accumulation
                        flag = p.device.graph_enabled()
                        p.device.EnableGraph(False)
                        self.gradAccumulation.append(
                            tensor.Tensor((acc,), p.device, p.dtype))
                        self.gradAccumulation[k].set_value(0.0)
                        p.device.EnableGraph(flag)
                    if corr:
                        self.fused_sparsification(glist,
                                                  self.gradAccumulation[k].data,
                                                  spars, topK)
                    else:
                        self.fused_sparsification(glist, None, spars, topK)
                    acc = 0
                    glist = []
            plist.append((p, g))
        if glist:
            k += 1
            if (corr and (not self.sparsInit)):
                # create a tensor for the gradient accumulation
                flag = p.device.graph_enabled()
                p.device.EnableGraph(False)
                self.gradAccumulation.append(
                    tensor.Tensor((acc,), p.device, p.dtype))
                self.gradAccumulation[k].set_value(0.0)
                p.device.EnableGraph(flag)
            if corr:
                self.fused_sparsification(glist, self.gradAccumulation[k].data,
                                          spars, topK)
            else:
                self.fused_sparsification(glist, None, spars, topK)
        self.wait()
        for p, g in plist:
            self.update(p, g)
        self.sparsInit = True
        self.opt.step()