def __init__(self, num_layers, input_dim, hidden_dim, output_dim):
        """MLP layers construction

        Paramters
        ---------
        num_layers: int
            The number of linear layers
        input_dim: int
            The dimensionality of input features
        hidden_dim: int
            The dimensionality of hidden units at ALL layers
        output_dim: int
            The number of classes for prediction
        """
        super(MLP, self).__init__()
        self.linear_or_not = True
        self.num_layers = num_layers
        self.output_dim = output_dim

        with self.name_scope():
            if num_layers < 1:
                raise ValueError("number of layers should be positive!")
            elif num_layers == 1:
                # Linear model
                self.linear = nn.Dense(output_dim, in_units=input_dim)
            else:
                self.linear_or_not = False
                self.linears = nn.Sequential()
                self.batch_norms = nn.Sequential()

                self.linears.add(nn.Dense(hidden_dim, in_units=input_dim))
                for layer in range(num_layers - 2):
                    self.linears.add(nn.Dense(hidden_dim, in_units=hidden_dim))
                self.linears.add(nn.Dense(output_dim, in_units=hidden_dim))

                for layer in range(num_layers - 1):
                    self.batch_norms.add(nn.BatchNorm(in_channels=hidden_dim))
MLP layers construction

        Paramters
        ---------
        num_layers: int
            The number of linear layers
        input_dim: int
            The dimensionality of input features
        hidden_dim: int
            The dimensionality of hidden units at ALL layers
        output_dim: int
            The number of classes for prediction
        
examples/mxnet/gin/gin.py