Vanishing seq_len in attention-based BiLSTM
I'm studying several implementations of self attention-based BiLSTM and I don't understand why in each of them the input and output size are different. In particular I refer to the following codes taken from different implementations:
Implementation 1 e 2
 def attnetwork(self, encoder_out, final_hidden):
        # encoder_out shape = (batch_size, seq_len, n_hidden)
        # final_hidden shape = (1, batch_size, n_hidden)
        hidden = final_hidden.squeeze(0)  
        attn_weights = torch.bmm(encoder_out, hidden.unsqueeze(2)).squeeze(2)
        soft_attn_weights = F.softmax(attn_weights, 1)
        new_hidden = torch.bmm(encoder_out.transpose(1,2), soft_attn_weights.unsqueeze(2)).squeeze(2)
        return new_hidden # shape = (batch_size, n_hidden) 

As you can see this implementation takes as input two vectors of dimension (batch_size, seq_len, n_hidden) and (1, batch_size, n_hidden), respectively, and returns a vector of dimensions (batch_size, n_hidden) . But where is the dimension relative to seq_len? I need to have an output vector equal to the input one (i.e (batch_size, seq_len, n_hidden)).
Another implementation where the input size does not match the output size:
def attention(self,H):
        M = torch.tanh(H) # Non-linear transformation size:(batch_size, hidden_dim, seq_len)
        a = F.softmax(torch.bmm(self.att_weight,M),dim=2) # a.Size : (batch_size,1,seq_len)
        a = torch.transpose(a,1,2) # (batch_size,seq_len,1)
        return torch.bmm(H,a) # (batch_size,hidden_dim,1)

Another implementation with the same problem:
def attention(self, rnn_out, state):
        merged_state = torch.cat([s for s in state],1)
        merged_state = merged_state.squeeze(0).unsqueeze(2)
        # (batch, seq_len, cell_size) * (batch, cell_size, 1) = (batch, seq_len, 1)
        weights = torch.bmm(rnn_out, merged_state)
        weights = torch.nn.functional.softmax(weights.squeeze(2)).unsqueeze(2)
        # (batch, cell_size, seq_len) * (batch, seq_len, 1) = (batch, cell_size, 1)
        return torch.bmm(torch.transpose(rnn_out, 1, 2), weights).squeeze(2)

How could one do to output a tensor of the same size as the input one without &quot;breaking&quot; the self attention mechanism?
Thank you!
EDIT:
The forward function I have to use is this:
def forward(self, x, x_len):
        x = nn.utils.rnn.pack_padded_sequence(x, x_len, batch_first=True)
        out1, (h_n, c_n) = self.lstm1(x)
        # out1 = (seq_len, batch, num_directions * hidden_size)
        # h_n = (num_layers * num_directions, batch, hidden_size)
        x, lengths = nn.utils.rnn.pad_packed_sequence(out1, batch_first=True)
        x, att1 = self.atten1(x, lengths)  # skip connect
        return x

the final x in return x I absolutely need it to have the shape (batch_size, seq_len, hidden_state) (obv also in another order so that a transpose is enough to fix it).

From my experience, what an attention-based model does is :

Calculate some relationships between decoder hidden states and
encoder outputs.
Take softmax to get the attention distribution.
Take a weighted sum of the encoder output to get the attention output.

And a Seq2seq model act like:

Pass your sequences e.g., an sentence through an encoder, get a encoder output and a hidden layer ,use this hidden layer to initialize your decoder.
Calculate your attention model by initial decoder hidden states and encoder outputs.
Feed your decoder with a start character e.g.,  ,get a single character as output and a new hidden state, then use your attention model to calculate a new output.
Feed your decoder with a single-character input (defined on your own, may depend on a teacher-force ratio) and attention output (may concatenate attention output with decoder hidden state or something else), then you get a new character.
Loop 4 to get your predict sequence character by character.

You may have a look at this (from page 59):
http://web.stanford.edu/class/cs224n/slides/cs224n-2021-lecture07-nmt.pdf
Every time a character sent to your LSTM decoder, it will output a new character, and its hidden states will update to another value. The key is, hidden states keep changing.
So a (batch_size,hidden_dim,1) attention output is reasonable, since you should give a character to your model, use a corresponding attention output to help predicting next character, then get a new hidden state and get a new attention output and so on. You may not give a whole sequence to your LSTM decoder, you can't use any hidden states under that circumstance.
