Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support Vararg Chain (Chain of Parallel) #2101

Closed
wants to merge 2 commits into from
Closed

Support Vararg Chain (Chain of Parallel) #2101

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
356d09d
Support Vararg Chain (Chain of Parallel)
cstjean Nov 7, 2022
0b19bcc
Documentation for Chain(Parallel(...))
cstjean Nov 7, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 20 additions & 11 deletions src/layers/basic.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,18 +48,21 @@ end

@functor Chain

(c::Chain)(x) = _applychain(c.layers, x)
(c::Chain)(inputs...) = _applychain(c.layers, inputs...)

@generated function _applychain(layers::Tuple{Vararg{<:Any,N}}, x) where {N}
symbols = vcat(:x, [gensym() for _ in 1:N])
calls = [:($(symbols[i+1]) = layers[$i]($(symbols[i]))) for i in 1:N]
Expr(:block, calls...)
@generated function _applychain(layers::Tuple{Vararg{<:Any,N}}, inputs...) where {N}
symbols = [gensym() for _ in 1:N]
calls = [:($(symbols[i]) = layers[$i]($(symbols[i-1]))) for i in 2:N]
Expr(:block,
:($(symbols[1]) = layers[1](inputs...)),
calls...)
end

_applychain(layers::NamedTuple, x) = _applychain(Tuple(layers), x)
_applychain(layers::NamedTuple, inputs...) = _applychain(Tuple(layers), inputs...)

function _applychain(layers::AbstractVector, x) # type-unstable path, helps compile times
for f in layers
function _applychain(layers::AbstractVector, inputs...) # type-unstable path, helps compile times
x = layers[1](inputs...)
for f in @view(layers[2:end])
x = f(x)
end
x
Expand Down Expand Up @@ -99,11 +102,11 @@ julia> activations(c, 1)
(2, 4, 64)
```
"""
activations(c::Chain, input) = _extraChain(Tuple(c.layers), input)
activations(c::Chain, inputs...) = _extraChain(Tuple(c.layers), inputs...)

# Calculates the forward results of each layer provided in a `Tuple` with `x` as model input.
function _extraChain(fs::Tuple, x)
res = first(fs)(x)
function _extraChain(fs::Tuple, inputs...)
res = first(fs)(inputs...)
return (res, _extraChain(Base.tail(fs), res)...)
end
_extraChain(::Tuple{}, x) = ()
Expand Down Expand Up @@ -490,6 +493,12 @@ julia> model2[:α](rand(10)) |> size

julia> model2[:β] == model2[2]
true

julia> model3 = Chain(Parallel(+, Dense(5 => 4), Embedding(15=>4)),
Dense(4 => 17));

julia> model3(randn(5, 10), rand(1:15, 10)) |> size
(17, 10)
```
"""
struct Parallel{F, T<:Union{Tuple, NamedTuple}}
Expand Down
7 changes: 6 additions & 1 deletion test/layers/basic.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
using Test, Random
import Flux: activations
import Flux: activations, OneHotArray, OneHotMatrix, OneHotVector, onehotbatch, params, Zygote

@testset "basic" begin
@testset "helpers" begin
Expand Down Expand Up @@ -216,6 +216,11 @@ import Flux: activations
@test size(Parallel(hcat, one = Dense(10, 10), two = identity)(input)) == (10, 4)
end

@testset "parallel chain" begin
inputs = (randn(2, 10), randn(3, 10))
@test size(Chain(Parallel(vcat, Dense(2, 5), identity), Dense(8, 4))(inputs...)) == (4, 10)
end

@testset "vararg input" begin
inputs = randn(10), randn(5), randn(4)
@test size(Parallel(+, Dense(10, 2), Dense(5, 2), Dense(4, 2))(inputs)) == (2,)
Expand Down