Merge pull request #467 from LuxDL/ap/repeated

Project.toml

@@ -1,7 +1,7 @@
 name = "Lux"
 uuid = "b2108857-7c20-44ae-9111-449ecde12c47"
 authors = ["Avik Pal <[email protected]> and contributors"]
-version = "0.5.11"
+version = "0.5.12"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

docs/src/api/Lux/layers.md

@@ -18,6 +18,7 @@ Chain
 PairwiseFusion
 Parallel
 SkipConnection
+RepeatedLayer
 ```
 
 ## Convolutional Layers

src/Lux.jl

@@ -57,7 +57,7 @@ end
 
 # Layers
 export cpu, gpu
-export Chain, Parallel, SkipConnection, PairwiseFusion, BranchLayer, Maxout
+export Chain, Parallel, SkipConnection, PairwiseFusion, BranchLayer, Maxout, RepeatedLayer
 export Bilinear, Dense, Embedding, Scale
 export Conv, ConvTranspose, CrossCor, MaxPool, MeanPool, GlobalMaxPool, GlobalMeanPool,
     AdaptiveMaxPool, AdaptiveMeanPool, Upsample, PixelShuffle

src/layers/containers.jl

@@ -574,3 +574,90 @@ Maxout(f::Function, n_alts::Int) = Maxout(ntuple(_ -> f(), n_alts)...)
 end
 
 Base.keys(m::Maxout) = Base.keys(getfield(m, :layers))
+
+"""
+    RepeatedLayer(model; repeats::Val = Val(10), input_injection::Val = Val(false))
+
+Iteratively applies `model` for `repeats` number of times. The initial input is passed
+into the model repeatedly if `input_injection = Val(true)`. This layer unrolls the
+computation, however, semantically this is same as:
+
+ 1. `input_injection = Val(false)`
+
+```julia
+res = x
+for i in 1:repeats
+    res, st = model(res, ps, st)
+end
+```
+
+ 2. `input_injection = Val(true)`
+
+```julia
+res = x
+for i in 1:repeats
+    res, st = model((res, x), ps, st)
+end
+```
+
+It is expected that `repeats` will be a reasonable number below `20`, beyond that compile
+times for gradients might be unreasonably high.
+
+## Arguments
+
+  - `model` must be an `AbstractExplicitLayer`
+
+## Keyword Arguments
+
+  - `repeats`: Number of times to apply the model
+  - `input_injection`: If `true`, then the input is passed to the model along with the
+    output
+
+## Inputs
+
+  - `x`: Input as described above
+
+## Returns
+
+  - Output is computed by as described above
+  - Updated state of the `model`
+
+## Parameters
+
+  - Parameters of `model`
+
+## States
+
+  - State of `model`
+"""
+struct RepeatedLayer{N, IJ, M <: AbstractExplicitLayer} <:
+       AbstractExplicitContainerLayer{(:model,)}
+    model::M
+end
+
+function LuxCore.display_name(model::RepeatedLayer{N, IJ}) where {N, IJ}
+    return "RepeatedLayer{repeats = $N, input_injection = $IJ}"
+end
+
+function RepeatedLayer(model::AbstractExplicitLayer; repeats::Val{N}=Val(10),
+        input_injection::Val{IJ}=Val(false)) where {N, IJ}
+    return RepeatedLayer{N, IJ, typeof(model)}(model)
+end
+
+(m::RepeatedLayer)(x, ps, st) = repeatedlayer(m, m.model, x, ps, st)
+
+@generated function repeatedlayer(::RepeatedLayer{N, IJ}, model, x, ps, st) where {N, IJ}
+    sts = ntuple(_ -> gensym("st"), N)
+    xs = ntuple(_ -> gensym("x"), N + IJ)
+    calls = []
+    IJ && push!(calls, :($(xs[1]) = x))
+    for i in 1:N
+        push!(calls,
+            :(($(xs[i + IJ]), $(sts[i])) = Lux.apply(model, $(IJ ? :(($(xs[i]), x)) : :x),
+                ps, $(i == 1 ? :st : sts[i - 1]))))
+    end
+    return quote
+        $(calls...)
+        return $(last(xs)), $(last(sts))
+    end
+end

test/layers/containers.jl

@@ -349,3 +349,24 @@ end
         @eval @test_gradients $__f $x $ps atol=1.0f-1 rtol=1.0f-1 gpu_testing=$ongpu
     end
 end
+
+@testset "$mode: Repeated" for (mode, aType, device, ongpu) in MODES
+    LAYERS = [Dense(2 => 2), Parallel(+, Dense(2 => 2), Dense(2 => 2)), Dense(2 => 2),
+        Parallel(+, Dense(2 => 2), Dense(2 => 2))]
+    REPEATS = [Val(4), Val(4), Val(4), Val(4)]
+    INJECTION = [Val(false), Val(true), Val(false), Val(true)]
+
+    @testset "repeats = $(repeats); input_injection = $(input_injection)" for (layer, repeats, input_injection) in zip(LAYERS,
+        REPEATS, INJECTION)
+        layer = RepeatedLayer(layer; repeats, input_injection)
+        display(layer)
+        ps, st = Lux.setup(rng, layer) .|> device
+        x = rand(rng, Float32, 2, 12) |> aType
+
+        @test size(layer(x, ps, st)[1]) == (2, 12)
+
+        @jet layer(x, ps, st)
+        __f = (x, ps) -> sum(first(layer(x, ps, st)))
+        @eval @test_gradients $__f $x $ps atol=1.0f-3 rtol=1.0f-3 gpu_testing=$ongpu
+    end
+end