# Path: nmt_chainer/additional_links/layer_normalization.py
# class LayerNormalizationLink(link.Link):
# 
#     """Layer normalization layer on outputs of linear functions.
#     This link implements a "layer normalization" layer
#     which normalizes the input units by statistics
#     that are computed along the second axis,
#     scales and shifts them.
#     Parameter initialization will be deferred until
#     the first forward data pass at which time the size will be determined.
#     Args:
#         size (int): Size of input units. If ``None``, parameter initialization
#             will be deferred until the first forward data pass at which time
#             the size will be determined.
#         eps (float): Epsilon value for numerical stability of normalization.
#         initial_gamma (~chainer.Initializer): Initializer for scaling vector.
#             If ``None``, then the vector is filled by 1.
#             If a scalar, the vector is filled by it.
#             If ``numpy.ndarray``, the vector is set by it.
#         initial_beta (~chainer.Initializer): Initializer for shifting vector.
#             If ``None``, then the vector is filled by 0.
#             If a scalar, the vector is filled by it.
#             If ``numpy.ndarray``, the vector is set by it.
#     Attributes:
#         gamma (~chainer.Parameter): Scaling parameter.
#         beta (~chainer.Parameter): Shifting parameter.
#         eps (float): Epsilon value for numerical stability.
#     See: `Layer Normalization <https://arxiv.org/abs/1607.06450>`_
#     """
# 
#     def __init__(self, size=None, eps=1e-6, initial_gamma=None,
#                  initial_beta=None):
#         super(LayerNormalizationLink, self).__init__()
#         self.add_uninitialized_param('gamma')
#         self.add_uninitialized_param('beta')
#         if initial_gamma is None:
#             initial_gamma = initializers.One()
#         self._gamma_initializer = initial_gamma
#         if initial_beta is None:
#             initial_beta = initializers.Zero()
#         self._beta_initializer = initial_beta
#         self.eps = eps
# 
#         if size is not None:
#             self._initialize_params(size)
# 
#     def _initialize_params(self, size):
#         self.add_param('gamma', (1,size))
#         initializers.init_weight(self.gamma.data, self._gamma_initializer)
#         self.add_param('beta', (1,size))
#         initializers.init_weight(self.beta.data, self._beta_initializer)
# 
#     def __call__(self, x):
#         """Apply layer normalization to given input.
#         Args:
#             x (~chainer.Variable): Batch vectors.
#                 Shape of this value must be `(batch_size, unit_size)`,
#                 e.g., the output of :func:`~chainer.functions.linear`.
#         Returns:
#             ~chainer.Variable: Output of the layer normalization.
#         """
# #         if self.gamma.data is None:
# #             self._initialize_params(x.size // x.shape[0])
# 
#         if self.has_uninitialized_params:
#             with cuda.get_device_from_id(self._device_id):
#                 self._initialize_params(x.size // x.shape[0])
# 
#         return layer_normalization(
#                 x, self.gamma, self.beta, self.eps)


import numpy
import six
import numpy
import six
import chainer
from chainer import cuda
from chainer import function
from chainer.utils import type_check
from chainer import cuda
from chainer.functions import lstm
from chainer.functions.array import concat
from chainer.functions.array import split_axis
from chainer import initializers
from chainer import link
from chainer.links.connection import linear
from chainer import variable
from nmt_chainer.additional_links.layer_normalization import LayerNormalizationLink as LayerNormalization


class LNStatelessLSTM(link.Chain):

    """Stateless LSTM layer.

    This is a fully-connected LSTM layer as a chain. Unlike the
    :func:`~chainer.functions.lstm` function, this chain holds upward and
    lateral connections as child links. This link doesn't keep cell and
    hidden states.

    Args:
        in_size (int or None): Dimension of input vectors. If ``None``,
            parameter initialization will be deferred until the first forward
            data pass at which time the size will be determined.
        out_size (int): Dimensionality of output vectors.

    Attributes:
        upward (chainer.links.Linear): Linear layer of upward connections.
        lateral (chainer.links.Linear): Linear layer of lateral connections.

    """

    def __init__(self, in_size, out_size,
                 lateral_init=None, upward_init=None,
                 bias_init=0, forget_bias_init=0):
        super(LNStatelessLSTM, self).__init__(
            upward=linear.Linear(in_size, 4 * out_size, initialW=0),
            lateral=linear.Linear(out_size, 4 * out_size,
                                  initialW=0, nobias=True),

# Path: nmt_chainer/additional_links/layer_normalization.py
# class LayerNormalizationLink(link.Link):
# 
#     """Layer normalization layer on outputs of linear functions.
#     This link implements a "layer normalization" layer
#     which normalizes the input units by statistics
#     that are computed along the second axis,
#     scales and shifts them.
#     Parameter initialization will be deferred until
#     the first forward data pass at which time the size will be determined.
#     Args:
#         size (int): Size of input units. If ``None``, parameter initialization
#             will be deferred until the first forward data pass at which time
#             the size will be determined.
#         eps (float): Epsilon value for numerical stability of normalization.
#         initial_gamma (~chainer.Initializer): Initializer for scaling vector.
#             If ``None``, then the vector is filled by 1.
#             If a scalar, the vector is filled by it.
#             If ``numpy.ndarray``, the vector is set by it.
#         initial_beta (~chainer.Initializer): Initializer for shifting vector.
#             If ``None``, then the vector is filled by 0.
#             If a scalar, the vector is filled by it.
#             If ``numpy.ndarray``, the vector is set by it.
#     Attributes:
#         gamma (~chainer.Parameter): Scaling parameter.
#         beta (~chainer.Parameter): Shifting parameter.
#         eps (float): Epsilon value for numerical stability.
#     See: `Layer Normalization <https://arxiv.org/abs/1607.06450>`_
#     """
# 
#     def __init__(self, size=None, eps=1e-6, initial_gamma=None,
#                  initial_beta=None):
#         super(LayerNormalizationLink, self).__init__()
#         self.add_uninitialized_param('gamma')
#         self.add_uninitialized_param('beta')
#         if initial_gamma is None:
#             initial_gamma = initializers.One()
#         self._gamma_initializer = initial_gamma
#         if initial_beta is None:
#             initial_beta = initializers.Zero()
#         self._beta_initializer = initial_beta
#         self.eps = eps
# 
#         if size is not None:
#             self._initialize_params(size)
# 
#     def _initialize_params(self, size):
#         self.add_param('gamma', (1,size))
#         initializers.init_weight(self.gamma.data, self._gamma_initializer)
#         self.add_param('beta', (1,size))
#         initializers.init_weight(self.beta.data, self._beta_initializer)
# 
#     def __call__(self, x):
#         """Apply layer normalization to given input.
#         Args:
#             x (~chainer.Variable): Batch vectors.
#                 Shape of this value must be `(batch_size, unit_size)`,
#                 e.g., the output of :func:`~chainer.functions.linear`.
#         Returns:
#             ~chainer.Variable: Output of the layer normalization.
#         """
# #         if self.gamma.data is None:
# #             self._initialize_params(x.size // x.shape[0])
# 
#         if self.has_uninitialized_params:
#             with cuda.get_device_from_id(self._device_id):
#                 self._initialize_params(x.size // x.shape[0])
# 
#         return layer_normalization(
#                 x, self.gamma, self.beta, self.eps)


# Path: nmt_chainer/additional_links/ln_lstm.py
import numpy
import six
import numpy
import six
import chainer
from chainer import cuda
from chainer import function
from chainer.utils import type_check
from chainer import cuda
from chainer.functions import lstm
from chainer.functions.array import concat
from chainer.functions.array import split_axis
from chainer import initializers
from chainer import link
from chainer.links.connection import linear
from chainer import variable
from nmt_chainer.additional_links.layer_normalization import LayerNormalizationLink as LayerNormalization



class LNStatelessLSTM(link.Chain):

    """Stateless LSTM layer.

    This is a fully-connected LSTM layer as a chain. Unlike the
    :func:`~chainer.functions.lstm` function, this chain holds upward and
    lateral connections as child links. This link doesn't keep cell and
    hidden states.

    Args:
        in_size (int or None): Dimension of input vectors. If ``None``,
            parameter initialization will be deferred until the first forward
            data pass at which time the size will be determined.
        out_size (int): Dimensionality of output vectors.

    Attributes:
        upward (chainer.links.Linear): Linear layer of upward connections.
        lateral (chainer.links.Linear): Linear layer of lateral connections.

    """

    def __init__(self, in_size, out_size,
                 lateral_init=None, upward_init=None,
                 bias_init=0, forget_bias_init=0):
        super(LNStatelessLSTM, self).__init__(
            upward=linear.Linear(in_size, 4 * out_size, initialW=0),
            lateral=linear.Linear(out_size, 4 * out_size,
                                  initialW=0, nobias=True),

            upward_ln = LayerNormalization(),