<jupyter_start><jupyter_text>This is an extendable implemtation of an LSTM, we'll use it as a starting point to investigate different structureTaken from here:  https://stackoverflow.com/questions/50168224/does-a-clean-and-extendable-lstm-implementation-exists-in-pytorch<jupyter_code>import math
import torch as th
import torch.nn as nn

class LSTMCell(nn.Module):

    def __init__(self, input_size, hidden_size, bias=True):
        super(LSTMCell, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.bias = bias
        self.i2h = nn.Linear(input_size, 4 * hidden_size, bias=bias)
        self.h2h = nn.Linear(hidden_size, 4 * hidden_size, bias=bias)
        self.reset_parameters()

    def reset_parameters(self):
        std = 1.0 / math.sqrt(self.hidden_size)
        for w in self.parameters():
            w.data.uniform_(-std, std)

    def forward(self, x, hidden):

        if hidden is None:
            hidden = self._init_hidden(x)
        
        h, c = hidden
        h = h.view(h.size(1), -1)
        c = c.view(c.size(1), -1)
        x = x.view(x.size(1), -1)

        print(x.shape)
        print(c.shape)
        print(h.shape)
        
        
        # Linear mappings
        preact = self.i2h(x) + self.h2h(h)

        # activations
        gates = preact[:, :3 * self.hidden_size].sigmoid()
        g_t = preact[:, 3 * self.hidden_size:].tanh()
        i_t = gates[:, :self.hidden_size]
        f_t = gates[:, self.hidden_size:2 * self.hidden_size]
        o_t = gates[:, -self.hidden_size:]

        c_t = th.mul(c, f_t) + th.mul(i_t, g_t)

        h_t = th.mul(o_t, c_t.tanh())

        h_t = h_t.view(1, h_t.size(0), -1)
        c_t = c_t.view(1, c_t.size(0), -1)
        return h_t, (h_t, c_t)

    
    
    @staticmethod
    def _init_hidden(input_):
        h = th.zeros_like(input_.view(1, input_.size(1), -1))
        c = th.zeros_like(input_.view(1, input_.size(1), -1))
        return h, c


class LSTM(nn.Module):

    def __init__(self, input_size, hidden_size, bias=True):
        super().__init__()
        self.lstm_cell = LSTMCell(input_size, hidden_size, bias)
        self.fc1 = nn.Linear(hidden_size, 1)

    def forward(self, input_, hidden=None):
        # input_ is of dimensionalty (1, time, input_size, ...)

        outputs = []
        for x in th.unbind(input_, dim=0):
            hidden = self.lstm_cell(x, hidden)
            outputs.append(hidden[0].clone())
        
        #output = self.fc1(outputs)
        
        return th.stack(outputs, dim=1)
lstm = LSTM(input_size = 4, hidden_size = 12, bias=False)

x = th.randn(10,1,4)
h = th.randn(1,12)
c = th.randn(1,12)

y = lstm.forward(x)
print(y.shape)
print(ht.shape)
print(ct.shape)
fc1 = nn.Linear(12,1)
inp = th.randn(12)

fc1(inp)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13907518531591953, 'fraction_numerical': 0.022590893046240734, 'mean_word_length': 2.365795724465558, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '549093', 'license': 'MIT', 'n_tokens_mistral': 1089, 'n_tokens_neox': 1059, 'n_words': 214, 'path': 'torch_lstm_mod/lstm.ipynb', 'repo_name': 'mmolnar0/sgillen_research', 'size': '2833'}