Adding layers and bidirectionality to custom LSTM cell in pytorch
I use a very custom LSTM-cell inspired by http://mlexplained.com/2019/02/15/building-an-lstm-from-scratch-in-pytorch-lstms-in-depth-part-1/. 

I use it to look at intermediate gating values. My question is, how would I expand this class to have an option for adding more layers and for adding bidirectionality? Should it be wrapped in a new class or added in the present one? 

class Dim(IntEnum):
    batch = 0
    seq = 1

class simpleLSTM(nn.Module):
    def __init__(self, input_sz: int, hidden_sz: int):
        super().__init__()
        self.input_size = input_sz
        self.hidden_size = hidden_sz
        # input gate
        self.W_ii = Parameter(torch.Tensor(input_sz, hidden_sz))
        self.W_hi = Parameter(torch.Tensor(hidden_sz, hidden_sz))
        self.b_i = Parameter(torch.Tensor(hidden_sz))
        # forget gate
        self.W_if = Parameter(torch.Tensor(input_sz, hidden_sz))
        self.W_hf = Parameter(torch.Tensor(hidden_sz, hidden_sz))
        self.b_f = Parameter(torch.Tensor(hidden_sz))
        # ???
        self.W_ig = Parameter(torch.Tensor(input_sz, hidden_sz))
        self.W_hg = Parameter(torch.Tensor(hidden_sz, hidden_sz))
        self.b_g = Parameter(torch.Tensor(hidden_sz))
        # output gate
        self.W_io = Parameter(torch.Tensor(input_sz, hidden_sz))
        self.W_ho = Parameter(torch.Tensor(hidden_sz, hidden_sz))
        self.b_o = Parameter(torch.Tensor(hidden_sz))

        self.init_weights()

        self.out = nn.Linear(hidden_sz, len(TRG.vocab))


    def init_weights(self):
        for p in self.parameters():
            if p.data.ndimension() &gt;= 2:
                nn.init.xavier_uniform_(p.data)
            else:
                nn.init.zeros_(p.data)

    def forward(self, x, init_states=None ):
        """Assumes x is of shape (batch, sequence, feature)"""
        seq_sz, bs, = x.size()
        hidden_seq = []
        prediction = []

        if init_states is None:
            h_t, c_t = torch.zeros(self.hidden_size).to(x.device), torch.zeros(self.hidden_size).to(x.device)
        else:
            h_t, c_t = init_states
        for t in range(seq_sz): # iterate over the time steps
            x_t = x[t, :].float()
            #LOOK HERE!!!
            i_t = torch.sigmoid(x_t @ self.W_ii + h_t @ self.W_hi + self.b_i)
            f_t = torch.sigmoid(x_t @ self.W_if + h_t @ self.W_hf + self.b_f)
            g_t = torch.tanh(x_t @ self.W_ig + h_t @ self.W_hg + self.b_g)
            o_t = torch.sigmoid(x_t @ self.W_io + h_t @ self.W_ho + self.b_o)
            c_t = f_t * c_t + i_t * g_t
            h_t = o_t * torch.tanh(c_t)

            hidden_seq.append(h_t.unsqueeze(Dim.batch))

            pred_t = self.out(h_t.unsqueeze(Dim.batch))
            #pred_t = F.softmax(pred_t)
            prediction.append(pred_t)


        hidden_seq = torch.cat(hidden_seq, dim=Dim.batch)
        prediction = torch.cat(prediction, dim=Dim.batch)
        # reshape from shape (sequence, batch, feature) to (batch, sequence, feature)
        hidden_seq = hidden_seq.transpose(Dim.batch, Dim.seq).contiguous()
        prediction = prediction.transpose(Dim.batch, Dim.seq).contiguous()


        return prediction, hidden_seq, (h_t, c_t)


I call it and train using the following as an example. 

lstm = simpleLSTM(1, 100)
hidden_size = lstm.hidden_size

optimizer = optim.Adam(lstm.parameters())


h_0, c_0 = (torch.zeros(hidden_size, requires_grad=True), 
            torch.zeros(hidden_size, requires_grad=True))

grads = []
h_t, c_t = h_0, c_0


N_EPOCHS = 10

for epoch in range(N_EPOCHS):
    epoch_loss = 0
    for i, batch in enumerate(train):

        optimizer.zero_grad()

        src, src_len = batch.src
        trg = batch.trg
        trg = trg.view(-1)


        predict, output, hidden_states = lstm(src)
        predict = predict.t().unsqueeze(1)
        predict= predict.view(-1, predict.shape[-1])

        loss = criterion(predict,trg)
        loss.backward()

        optimizer.step()
        epoch_loss += loss.item()

    print(epoch_loss)    


The easiest would be to create another module (say Bidirectional) and pass any cell you want to it.

Implementation itself is quite easy to do. Notice that I'm using concat operation for joining bi-directional output, you may want to specify other modes like summation etc.

Please read the comments in the code below, you may have to change it appropriately.

import torch


class Bidirectional(torch.nn.Module):
    def __init__(self, cell):
        super().__init__()
        self.cell = cell

    def __call__(self, x, init_states=None):
        prediction, hidden_seq, (h_t, c_t) = self.cell(x, init_states)
        backward_prediction, backward_hidden_seq, (
            backward_h_t,
            backward_c_t,
            # Assuming sequence is first dimension, otherwise change 0 appropriately
            # Reverses sequences so the LSTM cell acts on the reversed sequence
        ) = self.cell(torch.flip(x, (0,)), init_states)
        return (
            # Assuming you transpose so it has (batch, seq, features) dimensionality
            torch.cat((prediction, backward_prediction), 2),
            torch.cat((hidden_seq, backward_hidden_seq), 2),
            # Assuming it has (batch, features) dimensionality
            torch.cat((h_t, backward_ht), 1),
            torch.cat((c_t, backward_ct), 1),
        )


When it comes to multiple layers you could do something similiar in principle:

import torch


class Multilayer(torch.nn.Module):
    def __init__(self, *cells):
        super().__init__()
        self.cells = torch.nn.ModuleList(cells)

    def __call__(self, x, init_states=None):
        inputs = x
        for cell in self.cells:
            prediction, hidden_seq, (h_t, c_t) = cell(inputs, init_states)
            inputs = hidden_seq
        return prediction, hidden_seq, (h_t, c_t)


Please note you have to pass created cell objects into Multilayer e.g.:

# For three layers of LSTM, each needs features to be set up correctly
multilayer_LSTM = Multilayer(LSTM(), LSTM(), LSTM())


You may also pass classes instead of instances into constructor and create those inside Multilayer (so hidden_size matches automatically), but those ideas should get you started.
