# Source code for block_zoo.op.Concat2D

    # Copyright (c) Microsoft Corporation. All rights reserved.
    # Licensed under the MIT license.

    import torch
    import torch.nn as nn
    import logging

    from ..BaseLayer import BaseConf,BaseLayer
    from utils.DocInherit import DocInherit
    from utils.exceptions import ConfigurationError
    import copy

    class Concat2DConf(BaseConf):
        """ Configuration of Concat2D Layer

        Args:
            concat2D_axis(int): which axis to conduct concat2D, default is 1.
        """

        # init the args
        def __init__(self,**kwargs):
            super(Concat2DConf, self).__init__(**kwargs)

        # set default params
        @DocInherit
        def default(self):
            self.concat2D_axis = 1

        @DocInherit
        def declare(self):
            self.num_of_inputs = -1
            self.input_ranks = [2]

        @DocInherit
        def inference(self):
            self.output_dim = copy.deepcopy(self.input_dims[0])
            self.output_dim[-1] = 0
            self.output_dim[-1] += sum([input_dim[-1] for input_dim in self.input_dims])

            super(Concat2DConf, self).inference()

        @DocInherit
        def verify(self):
            super(Concat2DConf, self).verify()

            # to check if the ranks of all the inputs are equal
            rank_equal_flag = True
            for i in range(len(self.input_ranks)):
                if self.input_ranks[i] != self.input_ranks[0]:
                    rank_equal_flag = False
                    break
            if rank_equal_flag == False:
                raise ConfigurationError("For layer Concat2D, the ranks of each inputs should be equal!")

            # to check if the concat2D_axis is legal
            if self.concat2D_axis != 1:
                raise ConfigurationError("For layer Concat2D, the concat axis must be 1!")

    class Concat2D(nn.Module):
        """ Concat2D layer to merge sum of sequences(2D representation)

        Args:
            layer_conf (Concat2DConf): configuration of a layer
        """
        def __init__(self, layer_conf):
            super(Concat2D, self).__init__()
            self.layer_conf = layer_conf

            logging.warning("The length Concat2D layer returns is the length of first input")

        def forward(self, *args):
            """ process inputs

            Args:
                *args: (Tensor): string, string_len, string2, string2_len, ...
                    e.g. string (Tensor): [batch_size, dim], string_len (Tensor): [batch_size]

            Returns:
                Tensor: [batch_size, output_dim], [batch_size]

            """
            result = []
            for idx, input in enumerate(args):
                if idx % 2 == 0:
                    result.append(input)
            return torch.cat(result,self.layer_conf.concat2D_axis), args[1]