Python chainer.functions.split_axis() Examples

def _setup_slice(self, layer):
        if layer.slice_param.HasField('axis'):
            axis = layer.slice_param.axis
        elif layer.slice_param.HasField('slice_dim'):
            axis = layer.slice_param.slice_dim
        else:
            axis = 1

        if layer.slice_param.slice_point:
            indices_or_sections = list(layer.slice_param.slice_point)
        else:
            indices_or_sections = len(list(layer.top))

        self.forwards[layer.name] = _SingleArgumentFunction(
            functions.split_axis,
            indices_or_sections=indices_or_sections,
            axis=axis
        )

        self._add_layer(layer) 

def split_one_step_batch_input(xs):
    """Split one-step batch input.

    Args:
        xs (chainer.Variable, ndarray or tuple): One-step batched input. It
            should be either:
                - a variable whose first axis is the batch axis.
                - a tuple of such variables.

    Returns:
        list: Either a list of variables or a list of tuples of varialbes.
            The length of the list is the batch size of the input.
    """
    if isinstance(xs, tuple):
        return list(zip(*[split_one_step_batch_input(x) for x in xs]))
    else:
        return list(F.split_axis(xs, len(xs), axis=0)) 

def channelwise_inhibited(self, h):
        self.c = random.randint(0, 2)
        xp = cuda.get_array_module(h.data)
        num = h.data.shape[0]

        h = F.split_axis(h, 3, 1)
        c = F.reshape(h[self.c], (num, 16, 16))
        z = Variable(xp.zeros_like(c.data), 'AUTO')
        c = F.batch_matmul(c, z)
        c = F.reshape(c, (num, 1, 16, 16))
        hs = []
        for i, s in enumerate(h):
            if i == self.c:
                hs.append(c)
            else:
                hs.append(s)
        return F.concat(hs, 1) 

def __call__(self, prev_hg, prev_cg, prev_z, v, r, prev_u):
        v = self.broadcast_v(v)
        if r.shape[2] == 1:
            r = self.broadcast_r(r)

        lstm_input = cf.concat((prev_hg, v, r, prev_z), axis=1)
        gate_inputs = self.lstm(lstm_input)

        forget_gate_input, input_gate_input, tanh_input, output_gate_input = cf.split_axis(
            gate_inputs, 4, axis=1)

        forget_gate = cf.sigmoid(forget_gate_input)
        input_gate = cf.sigmoid(input_gate_input)
        next_c = forget_gate * prev_cg + input_gate * cf.tanh(tanh_input)
        output_gate = cf.sigmoid(output_gate_input)
        next_h = output_gate * cf.tanh(next_c)

        next_u = self.upsample_h(next_h) + prev_u

        return next_h, next_c, next_u 

def apply_to_seq(self, seq_list):
        mb_size = len(seq_list)
        mb_initial_cell, mb_initial_state = self.get_initial_states(mb_size)
        return self.nstep_lstm(mb_initial_cell, mb_initial_state, seq_list)


# class DoubleGRU(Chain):
#     def __init__(self, H, I):
#         log.info("using double GRU")
#         self.H1 = H/2
#         self.H2 = H - self.H1
#         super(DoubleGRU, self).__init__(
#             gru1 = faster_gru.GRU(self.H1, I),
#             gru2 = faster_gru.GRU(self.H2, self.H1)
#         )
#
#     def __call__(self, prev_state, inpt):
#         prev_state1, prev_state2 = F.split_axis(prev_state, (self.H1,), axis = 1)
#
#         prev_state1 = self.gru1(prev_state1, inpt)
#         prev_state2 = self.gru2(prev_state2, prev_state1)
#
#         return F.concat((prev_state1, prev_state2), axis = 1) 

def __call__(self, x):
        h = x
        for l in self.conv_layers:
            h = self.activation(l(h))

        # Advantage
        batch_size = x.shape[0]

        h = self.activation(self.main_stream(h))
        h_a, h_v = F.split_axis(h, 2, axis=-1)
        ya = F.reshape(self.a_stream(h_a),
                       (batch_size, self.n_actions, self.n_atoms))

        mean = F.sum(ya, axis=1, keepdims=True) / self.n_actions

        ya, mean = F.broadcast(ya, mean)
        ya -= mean

        # State value
        ys = F.reshape(self.v_stream(h_v), (batch_size, 1, self.n_atoms))
        ya, ys = F.broadcast(ya, ys)
        q = F.softmax(ya + ys, axis=2)

        return action_value.DistributionalDiscreteActionValue(q, self.z_values) 

def advance_state(self, previous_states, prev_y):
        current_mb_size = prev_y.data.shape[0]
        assert self.mb_size is None or current_mb_size <= self.mb_size

        if current_mb_size < len(previous_states[0].data):
            truncated_states = [None] * len(previous_states)
            for num_state in six.moves.range(len(previous_states)):
                truncated_states[num_state], _ = F.split_axis(
                    previous_states[num_state], (current_mb_size,), 0)
            previous_states = tuple(truncated_states)

        output_state = previous_states[-1]
        if self.decoder_chain.use_goto_attention:
            ci, attn = self.compute_ctxt(output_state, prev_y)
        else:
            ci, attn = self.compute_ctxt(output_state)
        concatenated = F.concat((prev_y, ci))

        new_states = self.decoder_chain.gru(previous_states, concatenated)
        return new_states, concatenated, attn 

def channelwise_inhibited(self, h):
        xp = cuda.get_array_module(h.data)
        num = h.data.shape[0]

        h = F.split_axis(h, 3, 1)
        c = F.reshape(h[self.c], (num, 16, 16))
        z = Variable(xp.zeros_like(c.data), 'AUTO')
        c = F.batch_matmul(c, z)
        c = F.reshape(c, (num, 1, 16, 16))
        hs = []
        for i, s in enumerate(h):
            if i == self.c:
                hs.append(c)
            else:
                hs.append(s)
        return F.concat(hs, 1) 

def faster_call2(self, h, x):
        r_z_h_x = self.W_r_z_h(x)

        r_z_h = self.U_r_z(h)

        r_x, z_x, h_x = split_axis(r_z_h_x, (self.n_units, self.n_units * 2), axis=1)
        assert r_x.data.shape[1] == self.n_units
        assert z_x.data.shape[1] == self.n_units
        assert h_x.data.shape[1] == self.n_units

        r_h, z_h = split_axis(r_z_h, (self.n_units,), axis=1)
#         r = sigmoid.sigmoid(r_x + r_h)
#         z = sigmoid.sigmoid(z_x + z_h)
#         h_bar = tanh.tanh(h_x + self.U(sigm_a_plus_b_by_h(r_x, r_h, h)))
#         h_new = (1 - z) * h + z * h_bar
#         return h_new

        return compute_output_GRU(z_x, z_h, h_x, h, self.U(sigm_a_plus_b_by_h_fast(r_x, r_h, h))) 

def forward(self, xs, ilens):
        '''BLSTM forward (the modified version)

        :param xs:
        :param ilens:
        :return:
        '''
        logging.info(self.__class__.__name__ + ' input lengths: ' + str(ilens))
        # need to move ilens to cpu
        ilens = cuda.to_cpu(ilens)
        hy, cy, ys = self.nblstm(None, None, xs)
        ys = self.l_last(F.vstack(ys))  # (sum _utt frame_utt) x dim
        xs = F.split_axis(ys, np.cumsum(ilens[:-1]), axis=0)
        del hy, cy

        # final tanh operation
        xs = F.split_axis(F.tanh(F.vstack(xs)), np.cumsum(ilens[:-1]), axis=0)

        # EDIT(hamaji): Unnecessary, as `force_tuple` is True by default.
        # # 1 utterance case, it becomes an array, so need to make a utt tuple
        # if not isinstance(xs, tuple):
        #     xs = [xs]

        return xs, ilens  # x: utt list of frame x dim 

def original(self, xs, ilens):
        '''BLSTM forward (the original implementation)

        :param xs:
        :param ilens:
        :return:
        '''
        logging.info(self.__class__.__name__ + ' input lengths: ' + str(ilens))
        # need to move ilens to cpu
        ilens = cuda.to_cpu(ilens)
        hy, cy, ys = self.nblstm(None, None, xs)
        ys = self.l_last(F.vstack(ys))  # (sum _utt frame_utt) x dim
        xs = F.split_axis(ys, np.cumsum(ilens[:-1]), axis=0)
        del hy, cy

        # final tanh operation
        xs = F.split_axis(F.tanh(F.vstack(xs)), np.cumsum(ilens[:-1]), axis=0)

        # 1 utterance case, it becomes an array, so need to make a utt tuple
        if not isinstance(xs, tuple):
            xs = [xs]

        return xs, ilens  # x: utt list of frame x dim 

def forward(self, xs, ilens):
        '''BLSTM forward (the modified version)

        :param xs:
        :param ilens:
        :return:
        '''
        logging.info(self.__class__.__name__ + ' input lengths: ' + str(ilens))
        # need to move ilens to cpu
        ilens = cuda.to_cpu(ilens)
        hy, cy, ys = self.nblstm(None, None, xs)
        ys = self.l_last(F.vstack(ys))  # (sum _utt frame_utt) x dim
        xs = F.split_axis(ys, np.cumsum(ilens[:-1]), 0)
        del hy, cy

        # final tanh operation
        xs = F.split_axis(F.tanh(F.vstack(xs)), np.cumsum(ilens[:-1]), 0)

        # EDIT(hamaji): Unnecessary, as `force_tuple` is True by default.
        # # 1 utterance case, it becomes an array, so need to make a utt tuple
        # if not isinstance(xs, tuple):
        #     xs = [xs]

        return xs, ilens  # x: utt list of frame x dim 

def test_type_hints(self):
        class Test():
            def forward(self, x: types.TyNdarray(np.float32, ('a', 'b'))):
                h = F.split_axis(x, 2, 1)
                return h

        model, forward_args = Test(), (np.zeros((10, 10)).astype(np.float32),)
        id2type = generate_id2type_from_forward(model, forward_args)

        self.assertEqual(str(id2type[1]), "class Test -> ndarray(float32, (10 (a), 10 (b))) -> (Variable(float32, (10 (a), 5 (b // 2))), Variable(float32, (10 (a), 5 (b // 2))))")	# FunctionDef forward (line 1)
        self.assertEqual(str(id2type[20]), "NoneType")	# Assign
        self.assertEqual(str(id2type[21]), "(Variable(float32, (10 (a), 5 (b // 2))), Variable(float32, (10 (a), 5 (b // 2))))")	# Name h (line 2)
        self.assertEqual(str(id2type[23]), "(Variable(float32, (10 (a), 5 (b // 2))), Variable(float32, (10 (a), 5 (b // 2))))")	# Call F.split_axis(x, 2, 1) (line 2)
        self.assertEqual(str(id2type[28]), "ndarray(float32, (10 (a), 10 (b)))")	# Name x (line 2)
        self.assertEqual(str(id2type[30]), "int")	# Constant 2 (line 2)
        self.assertEqual(str(id2type[31]), "int")	# Constant 1 (line 2)
        self.assertEqual(str(id2type[32]), "(Variable(float32, (10 (a), 5 (b // 2))), Variable(float32, (10 (a), 5 (b // 2))))")	# Return
        self.assertEqual(str(id2type[33]), "(Variable(float32, (10 (a), 5 (b // 2))), Variable(float32, (10 (a), 5 (b // 2))))")	# Name h (line 3) 

def __call__(self, prev_hg, prev_he, prev_ce, x, v, r, u):
        xu = cf.concat((x, u), axis=1)
        xu = self.downsample_xu(xu)
        v = self.broadcast_v(v)
        if r.shape[2] == 1:
            r = self.broadcast_r(r)

        lstm_input = cf.concat((prev_he, prev_hg, xu, v, r), axis=1)
        gate_inputs = self.lstm(lstm_input)

        if self.use_cuda_kernel:
            next_h, next_c = CoreFunction()(gate_inputs, prev_ce)
        else:
            forget_gate_input, input_gate_input, tanh_input, output_gate_input = cf.split_axis(
                gate_inputs, 4, axis=1)

            forget_gate = cf.sigmoid(forget_gate_input)
            input_gate = cf.sigmoid(input_gate_input)
            next_c = forget_gate * prev_ce + input_gate * cf.tanh(tanh_input)
            output_gate = cf.sigmoid(output_gate_input)
            next_h = output_gate * cf.tanh(next_c)

        return next_h, next_c 

def step(self, hx, cx, xs):
        """Batch of word tokens to word tokens and hidden LSTM states.

        Predict the next set of tokens given previous tokens.
        """
        # Concatenate all input captions and pass them through the model in a
        # single pass
        caption_lens = [len(x) for x in xs]
        caption_sections = np.cumsum(caption_lens[:-1])
        xs = F.concat(xs, axis=0)

        xs = self.embed_word(xs)
        xs = F.split_axis(xs, caption_sections, axis=0)
        hx, cx, ys = self.lstm(hx, cx, xs)

        ys = F.concat(ys, axis=0)
        ys = F.dropout(ys, self.dropout_ratio)
        ys = self.decode_caption(ys)
        return hx, cx, ys 

def __call__(self, x, h, c):
        hy = []
        cy = []
        for i, name in enumerate(self.x_amps.layer_names):
            hx_i = h[i]
            cx_i = c[i]
            gates = self.x_amps[name](x) + self.h_amps[name](hx_i)
            i_gate, f_gate, c_gate, o_gate = F.split_axis(gates, indices_or_sections=4, axis=1)
            i_gate = F.sigmoid(i_gate)
            f_gate = F.sigmoid(f_gate)
            c_gate = F.tanh(c_gate)
            o_gate = F.sigmoid(o_gate)
            cy_i = (f_gate * cx_i) + (i_gate * c_gate)
            hy_i = o_gate * F.sigmoid(cy_i)
            cy.append(cy_i)
            hy.append(hy_i)
            x = self.dropout(hy_i)
        return hy, cy 

def __call__(self, x):
        if self.downsample:
            y1 = self.dw_conv4(x)
            y1 = self.dw_bn4(y1)
            y1 = self.expand_conv5(y1)
            y1 = self.expand_bn5(y1)
            y1 = self.activ(y1)
            x2 = x
        else:
            y1, x2 = F.split_axis(x, indices_or_sections=2, axis=1)
        y2 = self.compress_conv1(x2)
        y2 = self.compress_bn1(y2)
        y2 = self.activ(y2)
        y2 = self.dw_conv2(y2)
        y2 = self.dw_bn2(y2)
        y2 = self.expand_conv3(y2)
        y2 = self.expand_bn3(y2)
        y2 = self.activ(y2)
        if self.use_se:
            y2 = self.se(y2)
        if self.use_residual and not self.downsample:
            y2 = y2 + x2
        x = F.concat((y1, y2), axis=1)
        x = self.c_shuffle(x)
        return x 

def inverse(self, y):
        scale_sqr = self.scale * self.scale
        batch, y_channels, y_height, y_width = y.shape
        assert (y_channels % scale_sqr == 0)
        x_channels = y_channels // scale_sqr
        x_height = y_height * self.scale
        x_width = y_width * self.scale

        x = F.transpose(y, axes=(0, 2, 3, 1))
        x = x.reshape(batch, y_height, y_width, scale_sqr, x_channels)
        d3_split_seq = F.split_axis(x, indices_or_sections=(x.shape[3] // self.scale), axis=3)
        d3_split_seq = [t.reshape(batch, y_height, x_width, x_channels) for t in d3_split_seq]
        x = F.stack(d3_split_seq, axis=0)
        x = F.transpose(F.swapaxes(x, axis1=0, axis2=1), axes=(0, 2, 1, 3, 4)).reshape(
            batch, x_height, x_width, x_channels)
        x = F.transpose(x, axes=(0, 3, 1, 2))
        return x 

def __call__(self, x):
        if self.downsample:
            y1 = self.shortcut_dconv(x)
            y1 = self.shortcut_conv(y1)
            x2 = x
        else:
            y1, x2 = F.split_axis(x, indices_or_sections=2, axis=1)
        y2 = self.conv1(x2)
        y2 = self.dconv(y2)
        y2 = self.conv2(y2)
        if self.use_se:
            y2 = self.se(y2)
        if self.use_residual and not self.downsample:
            y2 = y2 + x2
        x = F.concat((y1, y2), axis=1)
        x = self.c_shuffle(x)
        return x 

def forward(self, xs, hs=None, activation=None):
        if hs is not None:
            hx1, cx1, hx_emb, cx_emb = hs
        else:
            hx1 = cx1 = hx_emb = cx_emb = None
        # forward to LSTM layers
        hy_emb, cy_emb, ems = self.bi_lstm_emb(hx_emb, cx_emb, xs)
        hy1, cy1, ys = self.bi_lstm1(hx1, cx1, ems)
        # main branch
        ys_stack = F.vstack(ys)
        ys = self.linear1(ys_stack)
        if activation:
            ys = activation(ys)
        ilens = [x.shape[0] for x in xs]
        ys = F.split_axis(ys, np.cumsum(ilens[:-1]), axis=0)
        # embedding branch
        ems_stack = F.vstack(ems)
        ems = F.normalize(F.tanh(self.linear2(ems_stack)))
        ems = F.split_axis(ems, np.cumsum(ilens[:-1]), axis=0)

        if not isinstance(ys, tuple):
            ys = [ys]
            ems = [ems]
        return [hy1, cy1, hy_emb, cy_emb], ys, ems 

def __call__(self, s, xs):
        """Calculate all hidden states and cell states.
        Args:
            s  (~chainer.Variable or None): Initial (hidden & cell) states. If ``None``
                is specified zero-vector is used.
            xs (list of ~chianer.Variable): List of input sequences.
                Each element ``xs[i]`` is a :class:`chainer.Variable` holding
                a sequence.
        Return:
            (hy,cy): a pair of hidden and cell states at the end of the sequence,
            ys: a hidden state sequence at the last layer
        """
        if len(xs) > 1:
            sections = np.cumsum(np.array([len(x) for x in xs[:-1]], dtype=np.int32))
            xs = F.split_axis(self.embed(F.concat(xs, axis=0)), sections, axis=0)
        else:
            xs = [ self.embed(xs[0]) ]
        if s is not None:
            hy, cy, ys = self.lstm(s[0], s[1], xs)
        else:
            hy, cy, ys = self.lstm(None, None, xs)

        return (hy,cy), ys 

def __call__(self, s, xs):
        """Calculate all hidden states and cell states.
        Args:
            s  (~chainer.Variable or None): Initial (hidden & cell) states. If ``None``
                is specified zero-vector is used.
            xs (list of ~chianer.Variable): List of input sequences.
                Each element ``xs[i]`` is a :class:`chainer.Variable` holding
                a sequence.
        Return:
            (hy,cy): a pair of hidden and cell states at the end of the sequence,
            ys: a hidden state sequence at the last layer
        """
        if len(xs) > 1:
            sections = np.cumsum(np.array([len(x) for x in xs[:-1]], dtype=np.int32))
            xs = F.split_axis(self.embed(F.concat(xs, axis=0)), sections, axis=0)
        else:
            xs = [ self.embed(xs[0]) ]
        if s is not None:
            hy, cy, ys = self.lstm(s[0], s[1], xs)
        else:
            hy, cy, ys = self.lstm(None, None, xs)

        return (hy,cy), ys 

def __call__(self, s, xs):
        """Calculate all hidden states and cell states.
        Args:
            s  (~chainer.Variable or None): Initial (hidden & cell) states. If ``None``
                is specified zero-vector is used.
            xs (list of ~chianer.Variable): List of input sequences.
                Each element ``xs[i]`` is a :class:`chainer.Variable` holding
                a sequence.
        Return:
            (hy,cy): a pair of hidden and cell states at the end of the sequence,
            ys: a hidden state sequence at the last layer
        """
        if len(xs) > 1:
            sections = np.cumsum(np.array([len(x) for x in xs[:-1]], dtype=np.int32))
            xs = F.split_axis(self.embed(F.concat(xs, axis=0)), sections, axis=0)
        else:
            xs = [ self.embed(xs[0]) ]
        if s is not None:
            hy, cy, ys = self.lstm(s[0], s[1], xs)
        else:
            hy, cy, ys = self.lstm(None, None, xs)

        return (hy,cy), ys 

def original(self, xs, ilens):
        '''BLSTM forward (the original implementation)

        :param xs:
        :param ilens:
        :return:
        '''
        logging.info(self.__class__.__name__ + ' input lengths: ' + str(ilens))
        # need to move ilens to cpu
        ilens = cuda.to_cpu(ilens)
        hy, cy, ys = self.nblstm(None, None, xs)
        ys = self.l_last(F.vstack(ys))  # (sum _utt frame_utt) x dim
        xs = F.split_axis(ys, np.cumsum(ilens[:-1]), axis=0)
        del hy, cy

        # final tanh operation
        xs = F.split_axis(F.tanh(F.vstack(xs)), np.cumsum(ilens[:-1]), axis=0)

        # 1 utterance case, it becomes an array, so need to make a utt tuple
        if not isinstance(xs, tuple):
            xs = [xs]

        return xs, ilens  # x: utt list of frame x dim 

def split_batched_sequences(xs, sections):
    """Split concatenated sequences.

    Args:
        xs (chainer.Variable, ndarray or tuple): Concatenated sequences.
        sections (array-like): Sections as indices indicating start positions
            of sequences.

    Returns:
        list: List of sequences.
    """
    if isinstance(xs, tuple):
        return list(zip(*[split_batched_sequences(x, sections) for x in xs]))
    else:
        return list(F.split_axis(xs, sections, axis=0)) 

def sequence_embed(embed, xs, dropout=0.):
    """Efficient embedding function for variable-length sequences

    This output is equally to
    "return [F.dropout(embed(x), ratio=dropout) for x in xs]".
    However, calling the functions is one-shot and faster.

    Args:
        embed (callable): A :func:`~chainer.functions.embed_id` function
            or :class:`~chainer.links.EmbedID` link.
        xs (list of :class:`~chainer.Variable` or :class:`numpy.ndarray` or \
        :class:`cupy.ndarray`): i-th element in the list is an input variable,
            which is a :math:`(L_i, )`-shaped int array.
        dropout (float): Dropout ratio.

    Returns:
        list of ~chainer.Variable: Output variables. i-th element in the
        list is an output variable, which is a :math:`(L_i, N)`-shaped
        float array. :math:`(N)` is the number of dimensions of word embedding.

    """
    x_len = [len(x) for x in xs]
    x_section = np.cumsum(x_len[:-1])
    ex = embed(F.concat(xs, axis=0))
    ex = F.dropout(ex, ratio=dropout)
    exs = F.split_axis(ex, x_section, 0)
    return exs 

def sequence_embed(embed, xs, dropout=0.):
    """Efficient embedding function for variable-length sequences

    This output is equally to
    "return [F.dropout(embed(x), ratio=dropout) for x in xs]".
    However, calling the functions is one-shot and faster.

    Args:
        embed (callable): A :func:`~chainer.functions.embed_id` function
            or :class:`~chainer.links.EmbedID` link.
        xs (list of :class:`~chainer.Variable` or :class:`numpy.ndarray` or \
        :class:`cupy.ndarray`): i-th element in the list is an input variable,
            which is a :math:`(L_i, )`-shaped int array.
        dropout (float): Dropout ratio.

    Returns:
        list of ~chainer.Variable: Output variables. i-th element in the
        list is an output variable, which is a :math:`(L_i, N)`-shaped
        float array. :math:`(N)` is the number of dimensions of word embedding.

    """
    x_len = [len(x) for x in xs]
    x_section = numpy.cumsum(x_len[:-1])
    ex = embed(F.concat(xs, axis=0))
    ex = F.dropout(ex, ratio=dropout)
    exs = F.split_axis(ex, x_section, 0)
    return exs 

def __call__(self, x):
        xx = F.split_axis(x, self.in_channel_inds, axis=self.axis)
        out = [self[name_i](x_i) for x_i, name_i in zip(xx, self.layer_names)]
        x = F.concat(tuple(out), axis=self.axis)
        return x 

def sequence_embed(embed, xs, dropout=0.):
    """Efficient embedding function for variable-length sequences

    This output is equally to
    "return [F.dropout(embed(x), ratio=dropout) for x in xs]".
    However, calling the functions is one-shot and faster.

    Args:
        embed (callable): A :func:`~chainer.functions.embed_id` function
            or :class:`~chainer.links.EmbedID` link.
        xs (list of :class:`~chainer.Variable` or :class:`numpy.ndarray` or \
        :class:`cupy.ndarray`): i-th element in the list is an input variable,
            which is a :math:`(L_i, )`-shaped int array.
        dropout (float): Dropout ratio.

    Returns:
        list of ~chainer.Variable: Output variables. i-th element in the
        list is an output variable, which is a :math:`(L_i, N)`-shaped
        float array. :math:`(N)` is the number of dimensions of word embedding.

    """
    x_len = [len(x) for x in xs]
    x_section = numpy.cumsum(x_len[:-1])
    ex = embed(F.concat(xs, axis=0))
    ex = F.dropout(ex, ratio=dropout)
    exs = F.split_axis(ex, x_section, 0)
    return exs 

def inverse(self, x, _):
        x1, x2 = F.split_axis(x, indices_or_sections=2, axis=1)
        return x1, x2