# 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.
# =============================================================================

import math
from functools import wraps
from collections import OrderedDict

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


class LayerMeta(type):

    def init_wrapper(func):

        @wraps(func)
        def wrapper(self, *args, **kwargs):
            if len(args) == 0:
                return

            if isinstance(args[0], list):
                assert len(args) > 0 and isinstance(args[0][0], Tensor), (
                    'initialize function expects PlaceHolders or Tensors')
                dev = args[0][0].device
            else:
                assert len(args) > 0 and isinstance(args[0], Tensor), (
                    'initialize function expects PlaceHolders or Tensors')
                dev = args[0].device

            prev_state = dev.graph_enabled()
            dev.EnableGraph(False)
            func(self, *args, **kwargs)
            self._initialized = True
            dev.EnableGraph(prev_state)

        return wrapper

    def forward_wrapper(func):

        @wraps(func)
        def wrapper(self, *args, **kwargs):
            if not self._initialized:
                self.initialize(*args, **kwargs)
                self._initialized = True
            return func(self, *args, **kwargs)

        return wrapper

    def __new__(cls, name, bases, attr):
        if 'initialize' in attr:
            attr['initialize'] = LayerMeta.init_wrapper(attr['initialize'])
        if 'forward' in attr:
            attr['forward'] = LayerMeta.forward_wrapper(attr['forward'])

        return super(LayerMeta, cls).__new__(cls, name, bases, attr)


class Layer(object, metaclass=LayerMeta):

    sep = '.'

    def __init__(self):
        self.name = None
        self._initialized = False
        self._parent = None
        self._layers = dict()

    def initialize(self, *input):
        """ Initialize the layer

        This function will be called before the forward function if this
        layer hasn't been initialized. Those members that need to be
        initialized according to the input will be initialized in this
        function. e.g. parameters, states and handles.

        Args:
            *input: input args, should be consistent with the forward function
        """
        pass

    def forward(self, *input):
        """ Forward propagation

        Args:
            *input: input arguments consisting of only PyTensors
        Returns:
            PyTensor instance(s)
        """
        raise NotImplementedError

    def __call__(self, *args, **kwargs):
        return self.forward(*args, **kwargs)

    def get_params(self):
        """ Get parameters of this layer and all sublayers

        Returns:
            parameters(dict): A dictionary contains parameter names
            and values of this layer and all sublayers.
        """
        params = dict()
        sublayers = self._layers
        for name, sublayer in sublayers.items():
            if sublayer._initialized:
                params.update(sublayer.get_params())
        return params

    def set_params(self, parameters):
        """ Set parameters for this layer and all sublayers

        Args:
            parameters(dict): A dictionary contains parameter names
            and corresponding values. The value shoud be either a
            PyTensor or numpy ndarray
        """
        names = parameters.keys()
        sublayers = self._layers
        for name, sublayer in sublayers.items():
            if sublayer._initialized:
                if self._has_layer_param(sublayer, names):
                    sublayer.set_params(parameters)

    def get_states(self):
        """ Get states of this layer and all sublayers

        Returns:
            states(dict): A dictionary contains state names and values
            of this layer and all sublayers.
        """
        states = dict()
        sublayers = self._layers
        for name, sublayer in sublayers.items():
            if sublayer._initialized:
                states.update(sublayer.get_states())
        states.update(self.get_params())
        return states

    def set_states(self, states):
        """ Set states for this layer and all sublayers

        Args:
            states(dict): A dictionary contains state names and
            corresponding values. The value shoud be either a
            PyTensor or numpy ndarray
        """
        names = states.keys()
        sublayers = self._layers
        for name, sublayer in sublayers.items():
            if sublayer._initialized:
                if self._has_layer_param(sublayer, names):
                    sublayer.set_states(states)
        self.set_params(states)

    def dtype_check(self, *inputs):
        """ check if all input have same data type.

        Args:
            *inputs: input args consisting of only PyTensors
        """
        flag = inputs[0].device.graph_enabled()
        inputs[0].device.EnableGraph(False)

        x_dtype = inputs[0].dtype
        for inp in inputs:
            if inp.dtype != x_dtype:
                inp.to_type(x_dtype)

        inputs[0].device.EnableGraph(flag)

    def device_check(self, *inputs):
        """ Check if the devices of the input tensor are the same.

        Keep the device where each tensors is located the same as the
        first tensor. Copy data to the device of the first tensor if
        the device does not match.

        Args:
            *inputs: input args consisting of only PyTensors
        """
        # disabled the graph to prevent buffering data transfer operator
        x_device = inputs[0].device
        prev_state = x_device.graph_enabled()
        x_device.EnableGraph(False)
        x_dev_id = x_device.id()
        for var in inputs:
            if var.device.id() != x_dev_id:
                var.to_device(x_device)
        x_device.EnableGraph(prev_state)

    def _has_layer_param(self, layer, names):
        """ Determine whether names contains parameter names in the layer

        Args:
            layer(Layer): the layer instance
            names(list): the list of parameter names

        Returns:
            boolean: whether names contains parameter names in that layer
        """
        for name in names:
            if name.startswith(layer.name):
                return True
        return False

    def _get_name_prefix(self):
        """ Get the name prefix

        Returns:
            prefix(str): the layer or param name prefix
        """
        if self.name and self._parent:
            return self.name + Layer.sep
        else:
            return ''

    def __getattr__(self, name):
        if '_layers' in self.__dict__:
            layers = self.__dict__['_layers']
            if name in layers:
                return layers[name]
        raise AttributeError("'{}' object has no attribute '{}'".format(
            type(self).__name__, name))

    def __setattr__(self, name, value):
        if isinstance(value, Layer):
            # TODO: remove the attr from dict first
            self.__dict__['_layers'][name] = value
            value.__dict__['_parent'] = self
            value.name = self._get_name_prefix() + name
        else:
            object.__setattr__(self, name, value)
            if isinstance(value, Tensor) and value.is_dummy():
                # WARN: If tensors are initialized in __init__ function
                #       their names may be incorrect and should be reset
                value.name = self._get_name_prefix() + name
            elif name == 'name' and value:
                # WARN: can't reset the name after the initialization
                # update sublayer name
                for name, sublayer in self._layers.items():
                    sublayer.name = self._get_name_prefix() + name

    def __delattr__(self, name):
        if name in self._layers:
            del self._layers[name]
        else:
            object.__delattr__(self, name)

    def register_layers(self, *args):
        """ Register a list of sublayers.

        Can only be called once in each subclass.

        Args:
            *args: a list of sublayers or a dictionary that contains
            the name and the instance of each sublayer
        """
        if len(args) == 1 and isinstance(args[0], OrderedDict):
            items = args[0].items()
        else:
            items = [(v.__class__.__name__ + '_' + str(idx), v)
                     for idx, v in enumerate(args)]

        for name, value in items:
            if isinstance(value, Layer):
                self._layers[name] = value
                value.__dict__['_parent'] = self
                value.name = name


class Linear(Layer):
    """
    Generate a Linear operator
    """

    # TODO: replace current with
    #   def __init__(self, out_features, bias=True):
    def __init__(self, out_features, *args, bias=True, **kwargs):
        """
        Args:
            out_channels: int, the channel of output, also is the number of
                filters
            bias: bool
        """
        super(Linear, self).__init__()

        self.out_features = out_features

        # TODO: for backward compatibility, to remove
        if len(args) > 0:
            self.in_features = out_features
            self.out_features = args[0]
        if len(args) > 1:
            self.bias = args[1]
        else:
            self.bias = bias

    def initialize(self, x):
        self.in_features = x.shape[1]
        w_shape = (self.in_features, self.out_features)
        b_shape = (self.out_features,)

        self.W = Tensor(shape=w_shape,
                        dtype=x.dtype,
                        requires_grad=True,
                        stores_grad=True)
        std = math.sqrt(2.0 / (self.in_features + self.out_features))
        self.W.gaussian(0.0, std)

        if self.bias:
            self.b = Tensor(shape=b_shape,
                            dtype=x.dtype,
                            requires_grad=True,
                            stores_grad=True)
            self.b.set_value(0.0)
        else:
            self.b = None

    def forward(self, x):
        if self.b:
            self.device_check(x, self.W, self.b)
            self.dtype_check(x, self.W, self.b)
        else:
            self.device_check(x, self.W)
            self.dtype_check(x, self.W)

        assert x.shape[1] == self.W.shape[0], (
            "Linear layer expects input features size %d received %d" %
            (self.W.shape[0], x.shape[1]))

        y = autograd.matmul(x, self.W)
        if self.bias:
            y = autograd.add_bias(y, self.b, axis=0)
        return y

    def get_params(self):
        if self.bias:
            return {self.W.name: self.W, self.b.name: self.b}
        else:
            return {self.W.name: self.W}

    def set_params(self, parameters):
        self.W.copy_from(parameters[self.W.name])
        if self.bias:
            self.b.copy_from(parameters[self.b.name])


class Gemm(Layer):
    """
    Generate a Gemm operator
    Y = alpha * A' * B' + beta * C
    B is weight, C is bias
    """

    def __init__(self,
                 nb_kernels,
                 alpha=1.0,
                 beta=1.0,
                 transA=False,
                 transB=True,
                 bias=True,
                 bias_shape=None):
        """
        Args:
            nb_kernels: int, the channel of output, also is the number of
                filters
            alpha (float): Scalar multiplier for the product of input tensors A * B.
            beta (float): Scalar multiplier for input tensor C.
            ransA (bool): Whether A should be transposed
            transB (bool): Whether B should be transposed
            bias: bool
        """
        super(Gemm, self).__init__()
        self.nb_kernels = nb_kernels
        self.alpha = alpha
        self.beta = beta
        self.transA = 1 if transA else 0
        self.transB = 1 if transB else 0
        self.bias = bias
        self.bias_shape = bias_shape

    def initialize(self, x):
        if self.transA == 0:
            self.in_features = x.shape[-1]
        else:
            self.in_features = x.shape[0]

        if self.transB == 0:
            w_shape = (self.in_features, self.nb_kernels)
        else:
            w_shape = (self.nb_kernels, self.in_features)

        if self.bias_shape:
            b_shape = self.bias_shape
        else:
            b_shape = (1, self.nb_kernels)

        self.W = Tensor(shape=w_shape,
                        requires_grad=True,
                        stores_grad=True,
                        device=x.device)
        std = math.sqrt(2.0 / (self.in_features + self.nb_kernels))
        self.W.gaussian(0.0, std)

        if self.bias:
            self.b = Tensor(shape=b_shape,
                            requires_grad=True,
                            stores_grad=True,
                            device=x.device)
            self.b.set_value(0.0)
        else:
            self.b = None

    def forward(self, x):
        if self.b:
            self.device_check(x, self.W, self.b)
        else:
            self.device_check(x, self.W)

        if self.transA == 0:
            in_features = x.shape[-1]
        else:
            in_features = x.shape[0]

        if self.transB == 0:
            in_features_w = self.W.shape[0]
        else:
            in_features_w = self.W.shape[-1]

        assert in_features == in_features_w, (
            "Gemm layer expects input features size %d received %d" %
            (in_features_w, in_features))
        y = autograd.gemm(x, self.W, self.b, self.alpha, self.beta, self.transA,
                          self.transB)

        return y

    def get_params(self):
        if self.bias:
            return {self.W.name: self.W, self.b.name: self.b}
        else:
            return {self.W.name: self.W}

    def set_params(self, parameters):
        self.W.copy_from(parameters[self.W.name])
        if self.bias:
            self.b.copy_from(parameters[self.b.name])


class Embedding(Layer):
    """
    Generate an Embedding operator
    """

    def __init__(self, input_dim, output_dim, initializer="gaussian"):
        """init the Embedding operator
        Args:
            input_dim (int): the number of different words in the dictionary
            output_dim (int): the dimendion of a word after the embedding
            initializer (str, optional): weight initializer, can be [uniform, gaussian]. Defaults to "uniform".
        """
        super(Embedding, self).__init__()
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.initializer = initializer

    def initialize(self, x):
        w_shape = (self.input_dim, self.output_dim)
        self.W = Tensor(shape=w_shape,
                        requires_grad=True,
                        stores_grad=True,
                        device=x.device)
        if self.initializer == 'uniform':
            self.W.uniform(-1., 1.)
        else:
            self.W.gaussian(0., 1.)

    def from_pretrained(self, W, freeze=True):
        self.set_params({self.W.name: W})
        self.W.requires_grad = not freeze

    def forward(self, x):
        return autograd.embedding(x, self.W)

    def get_params(self):
        return {self.W.name: self.W}

    def set_params(self, parameters):
        self.W.copy_from(parameters[self.W.name])


class Conv2d(Layer):
    """
    Generate a Conv 2d operator
    """

    def __init__(self,
                 nb_kernels,
                 kernel_size,
                 *args,
                 stride=1,
                 padding=0,
                 dilation=1,
                 group=1,
                 bias=True,
                 pad_mode="NOTSET",
                 activation="NOTSET",
                 **kwargs):
        """
        Args:
            nb_kernels (int): the channel of output, also is the number of filters
            kernel_size (int or tuple): kernel size for two direction of each
                axis. For example, (2, 3), the first 2 means will add 2 at the
                beginning and also 2 at the end for its axis.and if a int is
                accepted, the kernel size will be initiated as (int, int)
            stride (int or tuple): stride, the logic is the same as kernel size.
            padding (int): tuple, list or None, padding, the logic is the same
                as kernel size. However, if you set pad_mode as "SAME_UPPER" or
                "SAME_LOWER" mode, you can set padding as None, and the padding
                will be computed automatically.
            dilation (int): only support 1
            group (int): group
            bias (bool): bias
            pad_mode (string): can be NOTSET, SAME_UPPER, or SAME_LOWER, where
                default value is NOTSET, which means explicit padding is used.
                SAME_UPPER or SAME_LOWER mean pad the input so that the output
                spatial size match the input. In case of odd number add the extra
                padding at the end for SAME_UPPER and at the beginning for SAME_LOWER.
            activation (string): can be NOTSET, RELU, where default value is NOTSET,
                which means there is no activation behind the conv2d layer.
                RELU means there is a ReLU behind current conv2d layer.
        """
        super(Conv2d, self).__init__()

        # the old code create the layer like: Conv2d(8, 16, 3)， or Conv2d(8, 16, 3, stride=1)
        # the following code block is for backward compatibility
        if len(args) > 0:
            nb_kernels = kernel_size
            kernel_size = args[0]
        if len(args) > 1:
            stride = args[1]
        if len(args) > 2:
            padding = args[2]

        self.nb_kernels = nb_kernels
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        self.dilation = dilation
        self.group = group
        self.bias = bias
        self.pad_mode = pad_mode
        self.activation = activation

        if isinstance(kernel_size, int):
            self.kernel_size = (kernel_size, kernel_size)
        elif isinstance(kernel_size, tuple):
            self.kernel_size = kernel_size
        else:
            raise TypeError("Wrong kernel_size type.")

        if isinstance(stride, int):
            self.stride = (stride, stride)
        elif isinstance(stride, tuple):
            self.stride = stride
        else:
            raise TypeError("Wrong stride type.")

        self.odd_padding = (0, 0, 0, 0)
        if isinstance(padding, int):
            self.padding = (padding, padding)
        elif isinstance(padding, tuple) or isinstance(padding, list):
            if len(padding) == 2:
                self.padding = padding
            elif len(padding) == 4:
                _h_mask = padding[0] - padding[1]
                _w_mask = padding[2] - padding[3]
                # the odd paddding is the value that cannot be handled by the tuple padding (w, h) mode
                # so we need to firstly handle the input, then use the nomal padding method.
                self.odd_padding = (max(_h_mask, 0), max(-_h_mask, 0),
                                    max(_w_mask, 0), max(-_w_mask, 0))
                self.padding = (
                    padding[0] - self.odd_padding[0],
                    padding[2] - self.odd_padding[2],
                )
            else:
                raise TypeError("Wrong padding value.")

        if dilation != 1 and list(dilation) != [1, 1]:
            raise ValueError("Not implemented yet")

        self.inner_params = {
            "cudnn_prefer": "fastest",
            "workspace_MB_limit": 1024,
        }
        # TODO valid value of inner_params check

        for kwarg in kwargs:
            if kwarg not in self.inner_params:
                raise TypeError("Keyword argument not understood:", kwarg)
            else:
                self.inner_params[kwarg] = kwargs[kwarg]

    def initialize(self, x):
        self.in_channels = x.shape[1]
        w_shape = (
            self.nb_kernels,
            int(self.in_channels / self.group),
            self.kernel_size[0],
            self.kernel_size[1],
        )

        self.W = Tensor(shape=w_shape,
                        requires_grad=True,
                        stores_grad=True,
                        device=x.device)
        # std = math.sqrt(
        # 2.0 / (self.in_channels * self.kernel_size[0] * self.kernel_size[1] +
        # self.nb_kernels))
        std = math.sqrt(
            2.0 / (w_shape[1] * self.kernel_size[0] * self.kernel_size[1] +
                   self.nb_kernels))
        self.W.gaussian(0.0, std)

        if self.bias:
            b_shape = (self.nb_kernels,)
            self.b = Tensor(shape=b_shape,
                            requires_grad=True,
                            stores_grad=True,
                            device=x.device)
            self.b.set_value(0.0)
        else:
            # to keep consistency when to do forward.
            self.b = None
            # Tensor(data=CTensor([]), requires_grad=False, stores_grad=False)

        # if same pad mode, re-compute the padding
        if self.pad_mode in ("SAME_UPPER", "SAME_LOWER"):
            self.padding, self.odd_padding = utils.get_padding_shape(
                self.pad_mode, x.shape[2:], self.kernel_size, self.stride)
            self.padding = [self.padding[0], self.padding[2]]

        _x = x
        if self.odd_padding != (0, 0, 0, 0):
            x_shape = list(x.data.shape())
            x_shape[2] += (self.odd_padding[0] + self.odd_padding[1])
            x_shape[3] += (self.odd_padding[2] + self.odd_padding[3])
            _x = Tensor(shape=x_shape, device=x.device)
            _x.set_value(0.0)

        if _x.device.id() == -1:
            if self.group != 1:
                raise ValueError("Not implemented yet")
            else:
                if not hasattr(self, "handle"):
                    self.handle = singa.ConvHandle(
                        _x.data,
                        self.kernel_size,
                        self.stride,
                        self.padding,
                        self.in_channels,
                        self.nb_kernels,
                        self.bias,
                        self.group,
                    )
        else:
            if not hasattr(self, "handle"):
                if _x.dtype == tensor.float16:
                    self.handle = singa.CudnnConvHandle(
                        _x.data,
                        self.kernel_size,
                        self.stride,
                        self.padding,
                        self.in_channels,
                        self.nb_kernels,
                        self.bias,
                        self.group,
                        1024*1024*1024,
                        "tensor_ops"
                    )
                else:
                    self.handle = singa.CudnnConvHandle(
                        _x.data,
                        self.kernel_size,
                        self.stride,
                        self.padding,
                        self.in_channels,
                        self.nb_kernels,
                        self.bias,
                        self.group,
                    )

    def forward(self, x):
        # sanitize the device of params/states, TODO: better to decorate forward()
        self.device_check(x, *[s for k, s in self.get_states().items()])
        self.dtype_check(x, *[s for k, s in self.get_states().items()])

        assert (self.group >= 1 and self.in_channels % self.group
                == 0), "please set reasonable group."

        assert (self.nb_kernels >= self.group and self.nb_kernels % self.group
                == 0), "nb_kernels and group dismatched."

        y = autograd.conv2d(self.handle, x, self.W, self.b, self.odd_padding)

        if self.activation != "NOTSET":
            if self.activation == "RELU":
                y = autograd.relu(y)

        return y

    def get_params(self):
        if self.bias:
            return {self.W.name: self.W, self.b.name: self.b}
        else:
            return {self.W.name: self.W}

    def set_params(self, parameters):
        self.W.copy_from(parameters[self.W.name])
        if self.bias:
            self.b.copy_from(parameters[self.b.name])


class SeparableConv2d(Layer):
    """
    Generate a Conv 2d operator
    """

    def __init__(self,
                 nb_kernels,
                 kernel_size,
                 *args,
                 stride=1,
                 padding=0,
                 bias=False):
        """
        Args:
            nb_kernels (int): the channel of output, also is the number of filters
            kernel_size (int or tuple): kernel size for two direction of each
                axis. For example, (2, 3), the first 2 means will add 2 at the
                beginning and also 2 at the end for its axis.and if a int is
                accepted, the kernel size will be initiated as (int, int)
            stride (int or tuple): stride, the logic is the same as kernel size.
            padding (int): tuple, list or None, padding, the logic is the same
                as kernel size. However, if you set pad_mode as "SAME_UPPER" or
                "SAME_LOWER" mode, you can set padding as None, and the padding
                will be computed automatically.
            bias (bool): bias
        """
        super(SeparableConv2d, self).__init__()

        # the following code block is for backward compatibility
        if len(args) > 0:
            nb_kernels = kernel_size
            kernel_size = args[0]
        if len(args) > 1:
            stride = args[1]
        if len(args) > 2:
            padding = args[2]

        self.nb_kernels = nb_kernels
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        self.bias = bias

    def initialize(self, x):
        self.in_channels = x.shape[1]
        self.depthwise_conv = Conv2d(
            self.in_channels,
            self.kernel_size,
            stride=self.stride,
            padding=self.padding,
            group=self.in_channels,
            bias=self.bias,
        )

        self.point_conv = Conv2d(self.nb_kernels, 1, bias=self.bias)

    def forward(self, x):
        y = self.depthwise_conv(x)
        y = self.point_conv(y)
        return y


class BatchNorm2d(Layer):
    """
    Generate a BatchNorm 2d operator
    """

    def __init__(self, *args, momentum=0.9):
        """
        Args:
            momentum (float): Factor used in computing the running mean and
                variance.
        """
        super(BatchNorm2d, self).__init__()

        if len(args) > 0:
            self.channels = args[0]
        if len(args) > 1:
            self.momentum = args[1]
        self.momentum = momentum
        assert 0 <= momentum <= 1.0, ("Illegal momentum")

    def initialize(self, x):
        self.channels = x.shape[1]
        param_shape = (self.channels,)

        self.scale = Tensor(shape=param_shape,
                            requires_grad=True,
                            stores_grad=True)
        self.scale.set_value(1.0)

        self.bias = Tensor(shape=param_shape,
                           requires_grad=True,
                           stores_grad=True)
        self.bias.set_value(0.0)

        self.running_mean = Tensor(shape=param_shape,
                                   requires_grad=False,
                                   stores_grad=False)
        self.running_mean.set_value(0.0)

        self.running_var = Tensor(shape=param_shape,
                                  requires_grad=False,
                                  stores_grad=False)
        self.running_var.set_value(1.0)

        if not hasattr(self, "handle"):
            if x.device.id() == -1:
                self.handle = singa.BatchNormHandle(self.momentum, x.data)
            else:
                self.handle = singa.CudnnBatchNormHandle(self.momentum, x.data)

    def forward(self, x):
        assert x.shape[1] == self.channels, (
            "number of channels dismatched. %d vs %d" %
            (x.shape[1], self.channels))

        self.device_check(x, self.scale, self.bias, self.running_mean,
                          self.running_var)
        self.dtype_check(x, self.scale, self.bias, self.running_mean,
                        self.running_var)

        y = autograd.batchnorm_2d(
            self.handle,
            x,
            self.scale,
            self.bias,
            self.running_mean,
            self.running_var,
        )
        return y

    def get_params(self):
        return {self.scale.name: self.scale, self.bias.name: self.bias}

    def set_params(self, parameters):
        self.scale.copy_from(parameters[self.scale.name])
        self.bias.copy_from(parameters[self.bias.name])

    def get_states(self):
        ret = self.get_params()
        ret[self.running_mean.name] = self.running_mean
        ret[self.running_var.name] = self.running_var
        return ret

    def set_states(self, states):
        self.set_params(states)
        self.running_mean.copy_from(states[self.running_mean.name])
        self.running_var.copy_from(states[self.running_var.name])


class Pooling2d(Layer):
    """
    Generate a Pooling 2d operator
    """

    def __init__(self,
                 kernel_size,
                 stride=None,
                 padding=0,
                 is_max=True,
                 pad_mode="NOTSET"):
        """
        Args:
            kernel_size (int or tuple): kernel size for two direction of each
                axis. For example, (2, 3), the first 2 means will add 2 at the
                beginning and also 2 at the end for its axis.and if a int is
                accepted, the kernel size will be initiated as (int, int)
            stride (int or tuple): stride, the logic is the same as kernel size.
            padding (int): tuple, list or None, padding, the logic is the same
                as kernel size. However, if you set pad_mode as "SAME_UPPER" or
                "SAME_LOWER" mode, you can set padding as None, and the padding
                will be computed automatically.
            is_max (bool): is max pooling or avg pooling
            pad_mode (string): can be NOTSET, SAME_UPPER, or SAME_LOWER, where
                default value is NOTSET, which means explicit padding is used.
                SAME_UPPER or SAME_LOWER mean pad the input so that the output
                spatial size match the input. In case of odd number add the extra
                padding at the end for SAME_UPPER and at the beginning for SAME_LOWER.
        """
        super(Pooling2d, self).__init__()

        if isinstance(kernel_size, int):
            self.kernel_size = (kernel_size, kernel_size)
        elif isinstance(kernel_size, tuple):
            self.kernel_size = kernel_size
        else:
            raise TypeError("Wrong kernel_size type.")

        if stride is None:
            self.stride = self.kernel_size
        elif isinstance(stride, int):
            self.stride = (stride, stride)
        elif isinstance(stride, tuple):
            self.stride = stride
            assert stride[0] > 0 or (kernel_size[0] == 1 and padding[0] == 0), (
                "stride[0]=0, but kernel_size[0]=%d, padding[0]=%d" %
                (kernel_size[0], padding[0]))
        else:
            raise TypeError("Wrong stride type.")

        self.odd_padding = (0, 0, 0, 0)
        if isinstance(padding, int):
            self.padding = (padding, padding)
        elif isinstance(padding, tuple) or isinstance(padding, list):
            if len(padding) == 2:
                self.padding = padding
            elif len(padding) == 4:
                _h_mask = padding[0] - padding[1]
                _w_mask = padding[2] - padding[3]
                # the odd paddding is the value that cannot be handled by the tuple padding (w, h) mode
                # so we need to firstly handle the input, then use the nomal padding method.
                self.odd_padding = (max(_h_mask, 0), max(-_h_mask, 0),
                                    max(_w_mask, 0), max(-_w_mask, 0))
                self.padding = (
                    padding[0] - self.odd_padding[0],
                    padding[2] - self.odd_padding[2],
                )
            else:
                raise TypeError("Wrong padding value.")

        self.is_max = is_max
        self.pad_mode = pad_mode

    def initialize(self, x):
        # if same pad mode, re-compute the padding
        if self.pad_mode in ("SAME_UPPER", "SAME_LOWER"):
            self.padding, self.odd_padding = utils.get_padding_shape(
                self.pad_mode, x.shape[2:], self.kernel_size, self.stride)

        # if same pad mode, re-compute the padding
        if self.pad_mode in ("SAME_UPPER", "SAME_LOWER"):
            self.padding, self.odd_padding = utils.get_padding_shape(
                self.pad_mode, x.shape[2:], self.kernel_size, self.stride)
            self.padding = [self.padding[0], self.padding[2]]

        _x = x
        if self.odd_padding != (0, 0, 0, 0):
            x_shape = list(x.data.shape())
            x_shape[2] += (self.odd_padding[0] + self.odd_padding[1])
            x_shape[3] += (self.odd_padding[2] + self.odd_padding[3])
            _x = Tensor(shape=x_shape, device=x.device)
            _x.set_value(0.0)

        if _x.device.id() == -1:
            self.handle = singa.PoolingHandle(
                _x.data,
                self.kernel_size,
                self.stride,
                self.padding,
                self.is_max,
            )
        else:
            self.handle = singa.CudnnPoolingHandle(
                _x.data,
                self.kernel_size,
                self.stride,
                self.padding,
                self.is_max,
            )

    def forward(self, x):
        y = autograd.pooling_2d(self.handle, x, self.odd_padding)
        return y


class MaxPool2d(Pooling2d):
    """
    Generate a Max Pooling 2d operator
    """

    def __init__(self, kernel_size, stride=None, padding=0, pad_mode="NOTSET"):
        """
        Args:
            kernel_size (int or tuple): kernel size for two direction of each
                axis. For example, (2, 3), the first 2 means will add 2 at the
                beginning and also 2 at the end for its axis.and if a int is
                accepted, the kernel size will be initiated as (int, int)
            stride (int or tuple): stride, the logic is the same as kernel size.
            padding (int): tuple, list or None, padding, the logic is the same
                as kernel size. However, if you set pad_mode as "SAME_UPPER" or
                "SAME_LOWER" mode, you can set padding as None, and the padding
                will be computed automatically.
            pad_mode (string): can be NOTSET, SAME_UPPER, or SAME_LOWER, where
                default value is NOTSET, which means explicit padding is used.
                SAME_UPPER or SAME_LOWER mean pad the input so that the output
                spatial size match the input. In case of odd number add the extra
                padding at the end for SAME_UPPER and at the beginning for SAME_LOWER.
        """
        super(MaxPool2d, self).__init__(kernel_size, stride, padding, True,
                                        pad_mode)


class AvgPool2d(Pooling2d):

    def __init__(self, kernel_size, stride=None, padding=0, pad_mode="NOTSET"):
        """
        Args:
            kernel_size (int or tuple): kernel size for two direction of each
                axis. For example, (2, 3), the first 2 means will add 2 at the
                beginning and also 2 at the end for its axis.and if a int is
                accepted, the kernel size will be initiated as (int, int)
            stride (int or tuple): stride, the logic is the same as kernel size.
            padding (int): tuple, list or None, padding, the logic is the same
                as kernel size. However, if you set pad_mode as "SAME_UPPER" or
                "SAME_LOWER" mode, you can set padding as None, and the padding
                will be computed automatically.
            pad_mode (string): can be NOTSET, SAME_UPPER, or SAME_LOWER, where
                default value is NOTSET, which means explicit padding is used.
                SAME_UPPER or SAME_LOWER mean pad the input so that the output
                spatial size match the input. In case of odd number add the extra
                padding at the end for SAME_UPPER and at the beginning for SAME_LOWER.
        """
        super(AvgPool2d, self).__init__(kernel_size, stride, padding, False,
                                        pad_mode)


class MaxPool1d(Pooling2d):
    """
    Generate a Max Pooling 1d operator
    """

    def __init__(self, kernel_size, stride=None, padding=0, pad_mode="NOTSET"):
        """
        Args:
            kernel_size (int or tuple): kernel size for two direction of each
                axis. For example, (2, 3), the first 2 means will add 2 at the
                beginning and also 2 at the end for its axis.and if a int is
                accepted, the kernel size will be initiated as (int, int)
            stride (int or tuple): stride, the logic is the same as kernel size.
            padding (int): tuple, list or None, padding, the logic is the same
                as kernel size. However, if you set pad_mode as "SAME_UPPER" or
                "SAME_LOWER" mode, you can set padding as None, and the padding
                will be computed automatically.
            pad_mode (string): can be NOTSET, SAME_UPPER, or SAME_LOWER, where
                default value is NOTSET, which means explicit padding is used.
                SAME_UPPER or SAME_LOWER mean pad the input so that the output
                spatial size match the input. In case of odd number add the extra
                padding at the end for SAME_UPPER and at the beginning for SAME_LOWER.
        """
        if stride is None:
            stride = kernel_size
        super(MaxPool1d, self).__init__((1, kernel_size), (1, stride),
                                        (0, padding), True, pad_mode)


class AvgPool1d(Pooling2d):
    """
    Generate a Avg Pooling 1d operator
    """

    def __init__(self, kernel_size, stride=None, padding=0, pad_mode="NOTSET"):
        """
        Args:
            kernel_size (int or tuple): kernel size for two direction of each
                axis. For example, (2, 3), the first 2 means will add 2 at the
                beginning and also 2 at the end for its axis.and if a int is
                accepted, the kernel size will be initiated as (int, int)
            stride (int or tuple): stride, the logic is the same as kernel size.
            padding (int): tuple, list or None, padding, the logic is the same
                as kernel size. However, if you set pad_mode as "SAME_UPPER" or
                "SAME_LOWER" mode, you can set padding as None, and the padding
                will be computed automatically.
            pad_mode (string): can be NOTSET, SAME_UPPER, or SAME_LOWER, where
                default value is NOTSET, which means explicit padding is used.
                SAME_UPPER or SAME_LOWER mean pad the input so that the output
                spatial size match the input. In case of odd number add the extra
                padding at the end for SAME_UPPER and at the beginning for SAME_LOWER.
        """
        if stride is None:
            stride = kernel_size
        super(AvgPool1d, self).__init__((1, kernel_size), (1, stride),
                                        (0, padding), False, pad_mode)


class RNN_Base(Layer):

    def step_forward(self,
                     x=None,
                     h=None,
                     c=None,
                     Wx=None,
                     Wh=None,
                     Bx=None,
                     Bh=None,
                     b=None):
        raise NotImplementedError


class RNN(RNN_Base):
    """
    Generate a RNN operator
    """

    def __init__(
        self,
        input_size,
        hidden_size,
        num_layers=1,
        nonlinearity="tanh",
        bias=True,
        batch_first=False,
        dropout=0,
        bidirectional=False,
    ):
        """
        Args:
            input_size (int):  The number of expected features in the input x
            hidden_size (int): The number of features in the hidden state h
            num_layers (int):  Number of recurrent layers. Default: 1
            nonlinearity (string): The non-linearity to use. Default: 'tanh'
            bias (bool):  If False, then the layer does not use bias weights.
                Default: True
            batch_first (bool):  If True, then the input and output tensors
                are provided as (batch, seq, feature). Default: False
            dropout (float): If non-zero, introduces a Dropout layer on the
                outputs of each RNN layer except the last layer, with dropout
                probability equal to dropout. Default: 0
            bidirectional (bool): If True, becomes a bidirectional RNN.
                Default: False
        """
        super(RNN, self).__init__()

        self.input_size = input_size
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.nonlinearity = nonlinearity
        self.bias = bias
        self.batch_first = batch_first
        self.dropout = dropout
        self.bidirectional = bidirectional

    def initialize(self, xs, h0):
        Wx_shape = (self.input_size, self.hidden_size)
        self.Wx = Tensor(shape=Wx_shape, requires_grad=True, stores_grad=True)
        self.Wx.gaussian(0.0, 1.0)

        Wh_shape = (self.hidden_size, self.hidden_size)
        self.Wh = Tensor(shape=Wh_shape, requires_grad=True, stores_grad=True)
        self.Wh.gaussian(0.0, 1.0)

        b_shape = (self.hidden_size,)
        self.b = Tensor(shape=b_shape, requires_grad=True, stores_grad=True)
        self.b.set_value(0.0)

    def forward(self, xs, h0):
        # xs: a tuple or list of input tensors
        if not isinstance(xs, tuple):
            xs = tuple(xs)
        inputs = xs + (h0,)
        self.device_check(*inputs)
        # self.device_check(inputs[0], *self.params)
        self.device_check(inputs[0], self.Wx, self.Wh, self.b)
        batchsize = xs[0].shape[0]
        out = []
        h = self.step_forward(xs[0], h0, self.Wx, self.Wh, self.b)
        out.append(h)
        for x in xs[1:]:
            assert x.shape[0] == batchsize
            h = self.step_forward(x, h, self.Wx, self.Wh, self.b)
            out.append(h)
        return out, h

    def step_forward(self, x, h, Wx, Wh, b):
        y2 = autograd.matmul(h, Wh)
        y1 = autograd.matmul(x, Wx)
        y = autograd.add(y2, y1)
        y = autograd.add_bias(y, b, axis=0)
        if self.nonlinearity == "tanh":
            y = autograd.tanh(y)
        elif self.nonlinearity == "relu":
            y = autograd.relu(y)
        else:
            raise ValueError
        return y

    def get_params(self):
        return {
            self.Wx.name: self.Wx,
            self.Wh.name: self.Wh,
            self.b.name: self.b
        }

    def set_params(self, parameters):
        self.Wx.copy_from(parameters[self.Wx.name])
        self.Wh.copy_from(parameters[self.Wh.name])
        self.b.copy_from(parameters[self.b.name])


class LSTM(RNN_Base):
    """
    Generate a LSTM operator
    """

    def __init__(
        self,
        input_size,
        hidden_size,
        nonlinearity="tanh",
        num_layers=1,
        bias=True,
        batch_first=False,
        dropout=0,
        bidirectional=False,
    ):
        """
        Args:
            input_size (int):  The number of expected features in the input x
            hidden_size (int): The number of features in the hidden state h
            num_layers (int):  Number of recurrent layers. Default: 1
            nonlinearity (string): The non-linearity to use. Default: 'tanh'
            bias (bool):  If False, then the layer does not use bias weights.
                Default: True
            batch_first (bool):  If True, then the input and output tensors
                are provided as (batch, seq, feature). Default: False
            dropout (float): If non-zero, introduces a Dropout layer on the
                outputs of each RNN layer except the last layer, with dropout
                probability equal to dropout. Default: 0
            bidirectional (bool): If True, becomes a bidirectional RNN.
                Default: False
        """
        super(LSTM, self).__init__()

        self.input_size = input_size
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.nonlinearity = nonlinearity
        self.bias = bias
        self.batch_first = batch_first
        self.dropout = dropout
        self.bidirectional = bidirectional

    def initialize(self, xs, h0_c0):
        # 1. Wx_i input,  Bx_i
        # 2. Wx_f forget, Bx_f
        # 3. Wx_o output, Bx_o
        # 4. Wx_g candidate, Bx_g
        Wx_shape = (self.input_size, self.hidden_size)
        self.Wx_i = Tensor(shape=Wx_shape, requires_grad=True, stores_grad=True)
        self.Wx_f = Tensor(shape=Wx_shape, requires_grad=True, stores_grad=True)
        self.Wx_o = Tensor(shape=Wx_shape, requires_grad=True, stores_grad=True)
        self.Wx_g = Tensor(shape=Wx_shape, requires_grad=True, stores_grad=True)

        Wh_shape = (self.hidden_size, self.hidden_size)
        self.Wh_i = Tensor(shape=Wh_shape, requires_grad=True, stores_grad=True)
        self.Wh_f = Tensor(shape=Wh_shape, requires_grad=True, stores_grad=True)
        self.Wh_o = Tensor(shape=Wh_shape, requires_grad=True, stores_grad=True)
        self.Wh_g = Tensor(shape=Wh_shape, requires_grad=True, stores_grad=True)
        [
            w.gaussian(0.0, 0.01) for w in [
                self.Wx_i, self.Wx_f, self.Wx_o, self.Wx_g, self.Wh_i,
                self.Wh_f, self.Wh_o, self.Wh_g
            ]
        ]

        Bx_shape = (self.hidden_size,)
        self.Bx_i = Tensor(shape=Bx_shape, requires_grad=True, stores_grad=True)
        self.Bx_f = Tensor(shape=Bx_shape, requires_grad=True, stores_grad=True)
        self.Bx_o = Tensor(shape=Bx_shape, requires_grad=True, stores_grad=True)
        self.Bx_g = Tensor(shape=Bx_shape, requires_grad=True, stores_grad=True)
        self.Bh_i = Tensor(shape=Bx_shape, requires_grad=True, stores_grad=True)
        self.Bh_f = Tensor(shape=Bx_shape, requires_grad=True, stores_grad=True)
        self.Bh_o = Tensor(shape=Bx_shape, requires_grad=True, stores_grad=True)
        self.Bh_g = Tensor(shape=Bx_shape, requires_grad=True, stores_grad=True)
        [
            b.set_value(0.0) for b in [
                self.Bx_i, self.Bx_f, self.Bx_o, self.Bx_g, self.Bh_i,
                self.Bh_f, self.Bh_o, self.Bh_g
            ]
        ]

    def forward(self, xs, h0_c0):
        # xs: a tuple or list of input tensors
        # h0_c0: a tuple of (h0, c0)
        h0, c0 = h0_c0
        if not isinstance(xs, list):
            xs = list(xs)
        inputs = xs + list((h0, c0))
        self.device_check(*inputs)
        self.device_check(inputs[0], *[s for k, s in self.get_states().items()])
        batchsize = xs[0].shape[0]
        out = []
        h, c = self.step_forward(xs[0], h0, c0)
        out.append(h)
        for x in xs[1:]:
            assert x.shape[0] == batchsize
            h, c = self.step_forward(x, h, c)
            out.append(h)
        return out, h, c

    def step_forward(self, x, h, c):
        # input
        y1 = autograd.matmul(x, self.Wx_i)
        y1 = autograd.add_bias(y1, self.Bx_i, axis=0)
        y2 = autograd.matmul(h, self.Wh_i)
        y2 = autograd.add_bias(y2, self.Bh_i, axis=0)
        i = autograd.add(y1, y2)
        i = autograd.sigmoid(i)

        # forget
        y1 = autograd.matmul(x, self.Wx_f)
        y1 = autograd.add_bias(y1, self.Bx_f, axis=0)
        y2 = autograd.matmul(h, self.Wh_f)
        y2 = autograd.add_bias(y2, self.Bh_f, axis=0)
        f = autograd.add(y1, y2)
        f = autograd.sigmoid(f)

        # output
        y1 = autograd.matmul(x, self.Wx_o)
        y1 = autograd.add_bias(y1, self.Bx_o, axis=0)
        y2 = autograd.matmul(h, self.Wh_o)
        y2 = autograd.add_bias(y2, self.Bh_o, axis=0)
        o = autograd.add(y1, y2)
        o = autograd.sigmoid(o)

        y1 = autograd.matmul(x, self.Wx_g)
        y1 = autograd.add_bias(y1, self.Bx_g, axis=0)
        y2 = autograd.matmul(h, self.Wh_g)
        y2 = autograd.add_bias(y2, self.Bh_g, axis=0)
        g = autograd.add(y1, y2)
        g = autograd.tanh(g)

        cout1 = autograd.mul(f, c)
        cout2 = autograd.mul(i, g)
        cout = autograd.add(cout1, cout2)

        hout = autograd.tanh(cout)
        hout = autograd.mul(o, hout)
        return hout, cout

    def get_params(self):
        ret = {}
        for w in [
                self.Wx_i, self.Wx_f, self.Wx_o, self.Wx_g, self.Wh_i,
                self.Wh_f, self.Wh_o, self.Wh_g
        ]:
            ret[w.name] = w

        for b in [
                self.Bx_i, self.Bx_f, self.Bx_o, self.Bx_g, self.Bh_i,
                self.Bh_f, self.Bh_o, self.Bh_g
        ]:
            ret[b.name] = b
        return ret

    def set_params(self, parameters):
        for w in [
                self.Wx_i, self.Wx_f, self.Wx_o, self.Wx_g, self.Wh_i,
                self.Wh_f, self.Wh_o, self.Wh_g
        ]:
            w.copy_from(parameters[w.name])

        for b in [
                self.Bx_i, self.Bx_f, self.Bx_o, self.Bx_g, self.Bh_i,
                self.Bh_f, self.Bh_o, self.Bh_g
        ]:
            b.copy_from(parameters[b.name])


''' layers without params or states
'''


class ReLU(Layer):
    """
    Generate a ReLU operator
    """

    def __init__(self):
        super(ReLU, self).__init__()

    def forward(self, x):
        return autograd.relu(x)


class Sigmoid(Layer):
    """
    Generate a ReLU operator
    """

    def __init__(self):
        super(Sigmoid, self).__init__()

    def forward(self, x):
        return autograd.sigmoid(x)


class Add(Layer):
    """
    Generate a Add operator
    """

    def __init__(self):
        super(Add, self).__init__()

    def forward(self, a, b):
        return autograd.add(a, b)


class Flatten(Layer):
    """
    Generate a Flatten operator
    """

    def __init__(self, axis=1):
        super(Flatten, self).__init__()
        self.axis = axis

    def forward(self, x):
        return autograd.flatten(x, self.axis)


class SoftMaxCrossEntropy(Layer):
    """
    Generate a SoftMaxCrossEntropy operator
    """

    def __init__(self):
        super(SoftMaxCrossEntropy, self).__init__()

    def forward(self, x, t):
        return autograd.softmax_cross_entropy(x, t)


class SoftMax(Layer):
    """
    Generate a SoftMax operator
    """

    def __init__(self):
        super(SoftMax, self).__init__()

    def forward(self, x):
        return autograd.softmax(x)


class MeanSquareError(Layer):
    """
    Generate a MeanSquareError operator
    """

    def __init__(self):
        super(MeanSquareError, self).__init__()

    def forward(self, x, t):
        return autograd.mse_loss(x, t)


class CrossEntropy(Layer):
    """
    Generate a CrossEntropy operator
    """

    def __init__(self):
        super(CrossEntropy, self).__init__()

    def forward(self, x, t):
        return autograd.cross_entropy(x, t)


class BinaryCrossEntropy(Layer):
    """
    Generate a BinaryCrossEntropy operator
    """

    def __init__(self):
        super(BinaryCrossEntropy, self).__init__()

    def forward(self, x, t):
        return autograd.binary_cross_entropy(x, t)


class Dropout(Layer):
    """
    Generate a Dropout operator
    """

    def __init__(self, ratio=0.5):
        super(Dropout, self).__init__()
        self.ratio = ratio

    def forward(self, x):
        return autograd.dropout(x, self.ratio)


class Cat(Layer):
    """
    Generate a Cat Operator
    """

    def __init__(self, axis=0):
        super(Cat, self).__init__()
        self.axis = axis

    def forward(self, xs):
        return autograd.cat(xs, self.axis)


class Reshape(Layer):
    """
    Generate a Reshape Operator
    """

    def __init__(self):
        super(Reshape, self).__init__()

    def forward(self, x, shape):
        return autograd.reshape(x, shape)


class CudnnRNN(Layer):
    """ `CudnnRNN` class implements with c++ backend and run the operation
          directly on cuDNN
        While `RNN` class implements with high level singa API
    """

    def __init__(self,
                 hidden_size,
                 activation="tanh",
                 num_layers=1,
                 bias=True,
                 batch_first=True,
                 dropout=0,
                 bidirectional=False,
                 rnn_mode="lstm",
                 use_mask=False,
                 return_sequences=True):
        """
            Args:
                hidden_size: hidden feature dim
                rnn_mode: accepted value: "vanilla", "tanh", "relu",  "lstm", "gru"
        """
        assert singa.USE_CUDA, "Not able to run without CUDA"
        assert num_layers > 0, "num layers should be > 0"
        assert 0 <= dropout < 1, "dropout shouldbe >=0 and <1"
        super(CudnnRNN, self).__init__()

        self.rnn_mode = rnn_mode
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.dropout = dropout
        self.bidirectional = 1 if bidirectional else 0
        self.return_sequences = return_sequences
        self.batch_first = batch_first
        self.use_mask = use_mask

        # GPU parameter
        # cudnn_rnn_mode: 0 - RNN RELU, 1 - RNN TANH, 2 - LSTM, 3 - GRU
        if self.rnn_mode == "lstm":
            self.cudnn_rnn_mode = 2
        elif self.rnn_mode == "vanilla" or self.rnn_mode == "tanh":
            self.cudnn_rnn_mode = 1
        elif self.rnn_mode == "relu":
            self.cudnn_rnn_mode = 0
        elif self.rnn_mode == "gru":
            self.cudnn_rnn_mode = 3

    def initialize(self, x, hx=None, cx=None, seq_lengths=None):
        if self.batch_first:
            x = x.transpose((1, 0, 2))
        self.input_size = x.shape[1]

        # GPU handle
        self.handle = singa.CudnnRNNHandle(x.data,
                                           self.hidden_size,
                                           mode=self.cudnn_rnn_mode,
                                           num_layers=self.num_layers,
                                           dropout=self.dropout,
                                           bidirectional=self.bidirectional)

        self.W = Tensor(shape=(self.handle.weights_size,),
                        requires_grad=True,
                        stores_grad=True,
                        device=x.device)

        k = 1 / self.hidden_size
        self.W.uniform(-math.sqrt(k), math.sqrt(k))

    def forward(self, x, hx=None, cx=None, seq_lengths=None):

        self.device_check(x, self.W)
        if self.batch_first:  # (bs,seq,data) -> (seq,bs,data)
            x = autograd.transpose(x, (1, 0, 2))

        batch_size = x.shape[1]
        directions = 2 if self.bidirectional else 1
        if hx == None:
            hx = Tensor(shape=(self.num_layers * directions, batch_size,
                               self.hidden_size),
                        requires_grad=False,
                        stores_grad=False,
                        device=x.device).set_value(0.0)
        if cx == None:
            cx = Tensor(shape=(self.num_layers * directions, batch_size,
                               self.hidden_size),
                        requires_grad=False,
                        stores_grad=False,
                        device=x.device).set_value(0.0)

        # outputs returned is list
        #   inputs has shape of {sequence length, batch size, feature size}
        if self.use_mask:
            assert type(seq_lengths) == Tensor, "wrong type for seq_lengths"
            y = autograd._RNN(self.handle,
                              return_sequences=self.return_sequences,
                              use_mask=self.use_mask,
                              seq_lengths=seq_lengths)(x, hx, cx, self.W)[0]
        else:
            y = autograd._RNN(
                self.handle,
                return_sequences=self.return_sequences,
            )(x, hx, cx, self.W)[0]
        if self.return_sequences and self.batch_first:
            # (seq, bs, hid) -> (bs, seq, hid)
            y = autograd.transpose(y, (1, 0, 2))
        return y

    def get_params(self):
        return {self.W.name: self.W}

    def set_params(self, parameters):
        self.set_attribute(self.W, parameters[self.W.name])


''' import autograd at the end to resolve circular import
'''
from singa import autograd