<commit_before>export recurrent
export LSTM, BiLSTM

function recurrent(f, x::Var, batchdims::Vector{Int}, h0::Var; rev=false)
    cumdims = Array{Int}(length(batchdims)+1)
    cumdims[1] = 1
    for i = 1:length(batchdims)
        cumdims[i+1] = cumdims[i] + batchdims[i]
    end
    perm = sortperm(batchdims, rev=true)
    h = h0
    hs = Var[]
    for t = 1:batchdims[perm[1]]
        xts = Var[]
        for p in perm
            t > batchdims[p] && break
            i = cumdims[p]
            i += rev ? batchdims[p]-t : t-1
            push!(xts, x[:,i:i])
        end
        xt = cat(2, xts...)
        if size(h,2) < size(xt,2)
            @assert size(h,2) == 1
            h = cat(2, ntuple(_ -> h, size(xt,2))...)
        elseif size(h,2) > size(xt,2)
            h = h[:,1:size(xt,2)]
        end
        xt = cat(1, xt, h)
        h = f(xt)
        push!(hs, h)
    end
    cat(2, hs...)
end

doc"""
    LSTM(::Type{T}, insize::Int, outsize::Int, [init_W=Uniform(0.001), init_U=Orthogonal()])

Long Short-Term Memory network.

```math
\begin{align*}
\mathbf{f}_{t} & =\sigma_{g}(W_{f}\mathbf{x}_{t}+U_{f}\mathbf{h}_{t-1}+\mathbf{b}_{f})\\
\mathbf{i}_{t} & =\sigma_{g}(W_{i}\mathbf{x}_{t}+U_{i}\mathbf{h}_{t-1}+\mathbf{b}_{i})\\
\mathbf{o}_{t} & =\sigma_{g}(W_{o}\mathbf{x}_{t}+U_{o}\mathbf{h}_{t-1}+\mathbf{b}_{o})\\
\mathbf{c}_{t} & =\mathbf{f}_{t}\odot\mathbf{c}_{t-1}+\mathbf{i}_{t}\odot\sigma_{c}(W_{c}\mathbf{x}_{t}+U_{c}\mathbf{h}_{t-1}+\mathbf{b}_{c})\\
\mathbf{h}_{t} & =\mathbf{o}_{t}\odot\sigma_{h}(\mathbf{c}_{t})
\end{align*}
```

* ``x_t \in R^{d}``: input vector to the LSTM block
* ``f_t \in R^{h}``: forget gate's activation vector
* ``i_t \in R^{h}``: input gate's activation vector
* ``o_t \in R^{h}``: output gate's activation vector
* ``h_t \in R^{h}``: output vector of the LSTM block
* ``c_t \in R^{h}``: cell state vector
* ``W \in R^{h \times d}``, ``U \in R^{h \times h}`` and ``b \in R^{h}``: weight matrices and bias vectors
* ``\sigma_g``: sigmoid function
* ``\sigma_c``: hyperbolic tangent function
* ``\sigma_h``: hyperbolic tangent function

# 👉 Example
```julia
T = Float32
x = Var(rand(T,100,10))
f = LSTM(T, 100, 100)
h = f(x)
```
"""
mutable struct LSTM
    WU::Var
    b::Var
    h0::Var
    c0::Var
end

function LSTM(::Type{T}, insize::Int, outsize::Int; init_W=Uniform(0.001), init_U=Orthogonal()) where T
    W = init_W(T, insize, 4outsize)
    U = init_U(T, insize, 4outsize)
    WU = cat(1, W, U)
    b = zeros(T, 4outsize)
    b[1:outsize] = ones(T, outsize) # forget gate initializes to 1
    h0 = zeros(T, outsize, 1)
    c0 = zeros(T, outsize, 1)
    LSTM(zerograd(WU), zerograd(b), zerograd(h0), zerograd(c0))
end
(lstm::LSTM)(x::Node, batchdims; name="") = Node(lstm, (x,batchdims), name)

function (lstm::LSTM)(x::Var, batchdims; rev=false)
    c = lstm.c0
    h = recurrent(x, batchdims, lstm.h0, rev=rev) do xt
        a = linear(xt, lstm.WU, lstm.b)
        n = size(a,1) ÷ 4
        f = sigmoid(a[1:n,:])
        i = sigmoid(a[n+1:2n,:])
        o = sigmoid(a[2n+1:3n,:])
        if size(c,2) < size(xt,2)<commit_msg>Fix recurrent bugs<commit_after>export recurrent
export LSTM, BiLSTM

function recurrent(f, x::Var, batchdims::Vector{Int}, h0::Var; rev=false)
    @assert sum(batchdims) == size(x,2)
    cumdims = Array{Int}(length(batchdims)+1)
    cumdims[1] = 1
    for i = 1:length(batchdims)
        cumdims[i+1] = cumdims[i] + batchdims[i]
    end
    perm = sortperm(batchdims, rev=true)
    h = h0
    hs = Array{Var}(size(x,2))
    for t = 1:batchdims[perm[1]]
        xts = Var[]
        for p in perm
            t > batchdims[p] && break
            i = cumdims[p]
            i += rev ? batchdims[p]-t : t-1
            push!(xts, x[:,i:i])
        end
        xt = cat(2, xts...)
        if size(h,2) < size(xt,2)
            @assert size(h,2) == 1
            h = cat(2, ntuple(_ -> h, size(xt,2))...)
        elseif size(h,2) > size(xt,2)
            h = h[:,1:size(xt,2)]
        end
        xt = cat(1, xt, h)
        h = f(xt)
        for j = 1:length(perm)
            p = perm[j]
            t > batchdims[p] && break
            i = cumdims[p]
            i += rev ? batchdims[p]-t : t-1
            hs[i] = h[:,j:j]
        end
    end
    cat(2, hs...)
end

doc"""
    LSTM(::Type{T}, insize::Int, outsize::Int, [init_W=Uniform(0.001), init_U=Orthogonal()])

Long Short-Term Memory network.

```math
\begin{align*}
\mathbf{f}_{t} & =\sigma_{g}(W_{f}\mathbf{x}_{t}+U_{f}\mathbf{h}_{t-1}+\mathbf{b}_{f})\\
\mathbf{i}_{t} & =\sigma_{g}(W_{i}\mathbf{x}_{t}+U_{i}\mathbf{h}_{t-1}+\mathbf{b}_{i})\\
\mathbf{o}_{t} & =\sigma_{g}(W_{o}\mathbf{x}_{t}+U_{o}\mathbf{h}_{t-1}+\mathbf{b}_{o})\\
\mathbf{c}_{t} & =\mathbf{f}_{t}\odot\mathbf{c}_{t-1}+\mathbf{i}_{t}\odot\sigma_{c}(W_{c}\mathbf{x}_{t}+U_{c}\mathbf{h}_{t-1}+\mathbf{b}_{c})\\
\mathbf{h}_{t} & =\mathbf{o}_{t}\odot\sigma_{h}(\mathbf{c}_{t})
\end{align*}
```

* ``x_t \in R^{d}``: input vector to the LSTM block
* ``f_t \in R^{h}``: forget gate's activation vector
* ``i_t \in R^{h}``: input gate's activation vector
* ``o_t \in R^{h}``: output gate's activation vector
* ``h_t \in R^{h}``: output vector of the LSTM block
* ``c_t \in R^{h}``: cell state vector
* ``W \in R^{h \times d}``, ``U \in R^{h \times h}`` and ``b \in R^{h}``: weight matrices and bias vectors
* ``\sigma_g``: sigmoid function
* ``\sigma_c``: hyperbolic tangent function
* ``\sigma_h``: hyperbolic tangent function

# 👉 Example
```julia
T = Float32
x = Var(rand(T,100,10))
f = LSTM(T, 100, 100)
h = f(x)
```
"""
mutable struct LSTM
    WU::Var
    b::Var
    h0::Var
    c0::Var
end

function LSTM(::Type{T}, insize::Int, outsize::Int; init_W=Xavier(), init_U=Orthogonal()) where T
    W = init_W(T, insize, 4outsize)
    U = init_U(T, insize, 4outsize)
    WU = cat(1, W, U)
    b = zeros(T, 4outsize)
    b[1:outsize] = ones(T, outsize) # forget gate initializes to 1
    h0 = zeros(T, outsize, 1)
    c0 = zeros(T, outsize, 1)
    LSTM(zerograd(WU), zerograd(b), zerograd(h0), zerograd(c0))
end
(lstm::LSTM)(x::Node, batchdims; name="") = Node(lstm, (x,batchdims), name)

function (lstm::LSTM)(x::Var, batchdims; rev=false)
    c = lstm.c0
    h = recurrent(x, batchdims, lstm.h0, rev=rev) do xt
        a = linear(xt, lstm.WU, lstm.b)
        n = size(a,1) ÷ 4
        f = sigmoid(a[1:n,:])
        i = sigmoid(a[n+1:2n,:])
        o = sigmoid(a[2n+1:3n,:])
        if size(c,2) < size(xt,2)
{'fraction_non_alphanumeric': 0.22095743021986589, 'fraction_numerical': 0.02338998908467176, 'mean_word_length': 2.9888059701492535, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 10, 'https://': 0, '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}}
{'dir': 'git-commits-cleaned', 'id': '1045674', 'n_tokens_mistral': 2791, 'n_tokens_neox': 2742, 'n_words': 581}