add init_hidden_state function (#101)

* add init_hidden_state function

* fix type and format of init_hidden_state

* use full function syntax for init_hidden_state

* rename init_hidden_state -> _init_hidden_state

* correct format detail

* add tests for _init_hidden_state

* fix _init_hidden_state tests

* fix format

* fix _init_hidden_state type instability

Co-authored-by: Avik Pal <[email protected]>

Co-authored-by: Avik Pal <[email protected]>

Project.toml

@@ -1,7 +1,7 @@
 name = "Lux"
 uuid = "b2108857-7c20-44ae-9111-449ecde12c47"
 authors = ["Avik Pal <[email protected]> and contributors"]
-version = "0.4.9"
+version = "0.4.10"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/layers/recurrent.jl

@@ -72,7 +72,7 @@ function (rnn::RNNCell)(x::AbstractMatrix, ps::Union{ComponentArray, NamedTuple}
                         st::NamedTuple)
     rng = replicate(st.rng)
     @set! st.rng = rng
-    hidden_state = rnn.init_state(rng, rnn.out_dims, size(x, 2))
+    hidden_state = _init_hidden_state(rng, rnn, x)
     return rnn((x, hidden_state), ps, st)
 end
 
@@ -206,8 +206,8 @@ function (lstm::LSTMCell)(x::AbstractMatrix, ps::Union{ComponentArray, NamedTupl
                           st::NamedTuple)
     rng = replicate(st.rng)
     @set! st.rng = rng
-    hidden_state = lstm.init_state(rng, lstm.out_dims, size(x, 2))
-    memory = lstm.init_state(rng, lstm.out_dims, size(x, 2))
+    hidden_state = _init_hidden_state(rng, lstm, x)
+    memory = _init_hidden_state(rng, lstm, x)
     return lstm((x, hidden_state, memory), ps, st)
 end
 
@@ -312,7 +312,7 @@ function (gru::GRUCell)(x::AbstractMatrix, ps::Union{ComponentArray, NamedTuple}
                         st::NamedTuple)
     rng = replicate(st.rng)
     @set! st.rng = rng
-    hidden_state = gru.init_state(rng, gru.out_dims, size(x, 2))
+    hidden_state = _init_hidden_state(rng, gru, x)
     return gru((x, hidden_state), ps, st)
 end
 

src/utils.jl

@@ -195,6 +195,15 @@ end
 @inline _gate(x::AbstractVector, h::Int, n::Int) = view(x, _gate(h, n))
 @inline _gate(x::AbstractMatrix, h::Int, n::Int) = view(x, _gate(h, n), :)
 
+@inline function _init_hidden_state(rng::AbstractRNG, rnn, x::AbstractMatrix)
+    return rnn.init_state(rng, rnn.out_dims, size(x, 2))
+end
+
+@inline function _init_hidden_state(rng::AbstractRNG, rnn,
+                                    x::Union{CUDA.StridedSubCuArray, CuArray})
+    return CuArray(rnn.init_state(rng, rnn.out_dims, size(x, 2)))
+end
+
 """
     multigate(x::AbstractArray, ::Val{N})
 

test/utils.jl

@@ -96,3 +96,15 @@ end
         @test_nowarn Optimisers.update!(st_opt, ps_c, ps_c)
     end
 end
+
+@testset "_init_hidden_state" begin
+    rnn = RNNCell(3 => 5; init_state=Lux.zeros32)
+    x = randn(rng, Float32, 3, 2, 2)
+    @test Lux._init_hidden_state(rng, rnn, view(x, :, 1, :)) == zeros(Float32, 5, 2)
+
+    if CUDA.functional()
+        x = x |> gpu
+        @test Lux._init_hidden_state(rng, rnn, view(x, :, 1, :)) ==
+              CUDA.zeros(Float32, 5, 2)
+    end
+end