PairwiseFusion layer, take 2

NEWS.md

@@ -1,5 +1,8 @@
 # Flux Release Notes
 
+## v0.13.4
+* Added [`PairwiseFusion` layer](https://github.com/FluxML/Flux.jl/pull/1983)
+
 ## v0.13
 * After a deprecations cycle, the datasets in `Flux.Data` have
 been removed in favour of MLDatasets.jl.

src/Flux.jl

@@ -16,7 +16,7 @@ export gradient
 # Pirate error to catch a common mistake. (Internal function `base` because overloading `update!` is more likely to give ambiguities.)
 Optimisers.base(dx::Zygote.Grads) = error("Optimisers.jl cannot be used with Zygote.jl's implicit gradients, `Params` & `Grads`")
 
-export Chain, Dense, Maxout, SkipConnection, Parallel,
+export Chain, Dense, Maxout, SkipConnection, Parallel, PairwiseFusion,
        RNN, LSTM, GRU, GRUv3,
        SamePad, Conv, CrossCor, ConvTranspose, DepthwiseConv,
        AdaptiveMaxPool, AdaptiveMeanPool, GlobalMaxPool, GlobalMeanPool, MaxPool, MeanPool,

src/layers/basic.jl

@@ -67,7 +67,7 @@ end
 
 Base.getindex(c::Chain, i::AbstractArray) = Chain(c.layers[i])
 Base.getindex(c::Chain{<:NamedTuple}, i::AbstractArray) =
-  Chain(NamedTuple{Base.keys(c)[i]}(Tuple(c.layers)[i]))
+  Chain(NamedTuple{keys(c)[i]}(Tuple(c.layers)[i]))
 function Base.show(io::IO, c::Chain)
   print(io, "Chain(")
   _show_layers(io, c.layers)
@@ -487,7 +487,7 @@ end
 Parallel(connection, layers...) = Parallel(connection, layers)
 function Parallel(connection; kw...)
   layers = NamedTuple(kw)
-  if :layers in Base.keys(layers) || :connection in Base.keys(layers)
+  if :layers in keys(layers) || :connection in keys(layers)
     throw(ArgumentError("a Parallel layer cannot have a named sub-layer called `connection` or `layers`"))
   end
   isempty(layers) && return Parallel(connection, ())
@@ -510,16 +510,114 @@ end
 Base.getindex(m::Parallel, i) = m.layers[i]
 Base.getindex(m::Parallel, i::AbstractVector) = Parallel(m.connection, m.layers[i])
 Base.getindex(m::Parallel{<:Any, <:NamedTuple}, i::AbstractVector) =
-  Parallel(m.connection, NamedTuple{Base.keys(m)[i]}(Tuple(m.layers)[i]))
+  Parallel(m.connection, NamedTuple{keys(m)[i]}(Tuple(m.layers)[i]))
 
-Base.keys(m::Parallel) = Base.keys(getfield(m, :layers))
+Base.keys(m::Parallel) = keys(getfield(m, :layers))
 
 function Base.show(io::IO, m::Parallel)
   print(io, "Parallel(", m.connection, ", ")
   _show_layers(io, m.layers)
   print(io, ")")
 end
 
+"""
+    PairwiseFusion(connection, layers...)
+
+## Arguments
+
+- `connection`: A function taking 2 inputs and combining them into a single output 
+- `layers`: The layers whose outputs are combined
+
+## Inputs
+
+This layer behaves differently based on input type:
+
+1. Input `x` is a tuple of length `N`. Then `layers` must be a tuple of length `N`. 
+
+With 3 layers, it takes 3 inputs and returns 3 outputs.
+`(y1, y2, y3) = PairwiseFusion(connection, layer1, layer2, layer3)((x1, x2, x3))`
+may be drawn as:
+```
+x1 → layer1 → y1 ↘
+                  connection → layer2 → y2 ↘
+              x2 ↗                          connection → layer3 → y3
+                                        x3 ↗
+```
+
+In code:
+```julia
+y1 = layer1(x1)
+y2 = layer2(connection(x2, y1))
+y3 = layer3(connection(x3, y2))
+```
+
+2. Any other kind of input:
+
+```julia
+y₁ = x
+for i in 1:N
+    yᵢ₊₁ = connection(x, layers[i](yᵢ))
+end
+```
+
+## Returns
+
+A tuple of length `N` with the output of each fusion ((`y1`, `y2`, ..., `yN`) in the example above).
+"""
+struct PairwiseFusion{F, T <: NamedTuple}
+  connection::F
+  layers::T
+end
+
+function PairwiseFusion(connection, layers...)
+  names = ntuple(i -> Symbol("layer_$i"), length(layers))
+  return PairwiseFusion(connection, NamedTuple{names}(layers))
+end
+
+function _pairwise_check(lx, N, T)
+  if T <: Tuple && lx != N
+    throw(ArgumentError("PairwiseFusion with $N sub-layers can take one input or $N inputs, but got $lx inputs"))
+  end
+end
+ChainRulesCore.@non_differentiable _pairwise_check(lx, N, T)
+
+function (m::PairwiseFusion)(x::T) where {T}
+  lx = length(x)
+  N = length(m.layers)
+  _pairwise_check(lx, N, T)
+  applypairwisefusion(m.layers, m.connection, x)
+end
+
+@generated function applypairwisefusion(layers::NamedTuple{names}, connection, x::T) where {names, T}
+  N = length(names)
+  y_symbols = [gensym() for _ in 1:(N + 1)]
+  getinput(i) = T <: Tuple ? :(x[$i]) : :x
+  calls = [:($(y_symbols[N + 1]) = $(getinput(1)))]
+  for i in 1:N - 1
+    push!(calls, quote
+      $(y_symbols[i]) = layers[$i]($(y_symbols[N + 1]))
+      $(y_symbols[N + 1]) = connection($(y_symbols[i]), $(getinput(i + 1)))
+    end)
+  end
+  push!(calls, :($(y_symbols[N]) = layers[$N]($(y_symbols[N + 1]))))
+  push!(calls, :(return tuple($(Tuple(y_symbols[1:N])...))))
+  return Expr(:block, calls...)
+end
+
+@functor PairwiseFusion
+
+Base.getindex(m::PairwiseFusion, i) = m.layers[i]
+Base.getindex(m::PairwiseFusion{<:Any, <:NamedTuple}, i::AbstractVector) =
+  PairwiseFusion(m.connection, NamedTuple{keys(m)[i]}(Tuple(m.layers)[i]))
+
+Base.keys(m::PairwiseFusion) = keys(getfield(m, :layers))
+
+function Base.show(io::IO, m::PairwiseFusion)
+  print(io, "PairwiseFusion(", m.connection, ", ")
+  _show_layers(io, m.layers)
+  print(io, ")")
+end
+
 """
     Embedding(in => out; init=randn)
 
@@ -556,7 +654,7 @@ end
 @functor Embedding
 
 Embedding((in, out)::Pair{<:Integer, <:Integer}; init = randn32) = Embedding(init(out, in))
-  
+
 (m::Embedding)(x::Integer) = m.weight[:, x]
 (m::Embedding)(x::AbstractVector) = NNlib.gather(m.weight, x)
 (m::Embedding)(x::AbstractArray) = reshape(m(vec(x)), :, size(x)...)
@@ -565,7 +663,7 @@ function (m::Embedding)(x::Union{OneHotVector{T,L}, OneHotMatrix{T,L}}) where {T
   size(m.weight, 2) == L || throw(DimensionMismatch("Matrix column must correspond with OneHot size: $(size(m.weight, 2)) != $L"))
   return m(onecold(x))
 end
- 
+
 function Base.show(io::IO, m::Embedding)
   print(io, "Embedding(", size(m.weight, 2), " => ", size(m.weight, 1), ")")
 end

src/layers/show.jl

@@ -1,6 +1,6 @@
 
 for T in [
-    :Chain, :Parallel, :SkipConnection, :Recur, :Maxout  # container types
+    :Chain, :Parallel, :SkipConnection, :Recur, :Maxout, :PairwiseFusion  # container types
   ]
   @eval function Base.show(io::IO, m::MIME"text/plain", x::$T)
     if get(io, :typeinfo, nothing) === nothing  # e.g. top level in REPL
@@ -53,6 +53,7 @@ _show_children(x) = trainable(x)  # except for layers which hide their Tuple:
 _show_children(c::Chain) = c.layers
 _show_children(m::Maxout) = m.layers
 _show_children(p::Parallel) = (p.connection, p.layers...)
+_show_children(f::PairwiseFusion) = (f.connection, f.layers...)
 
 for T in [
     :Conv, :ConvTranspose, :CrossCor, :Dense, :Scale, :Bilinear, :Embedding,

test/layers/basic.jl

@@ -350,3 +350,19 @@ end
   @test Flux.destructure(m1)[2](z1)[1].weight ≈ Flux.destructure(m1v)[2](z1)[1].weight
   # Note that Flux.destructure(m1v)[2](z) has a Chain{Tuple}, as does m1v[1:2]
 end
+
+@testset "PairwiseFusion" begin
+  x = (rand(1, 10), rand(30, 10))
+  layer = PairwiseFusion(+,  Dense(1, 30),  Dense(30, 10))
+  y = layer(x)
+  @test length(y) == 2
+  @test size(y[1]) == (30, 10)
+  @test size(y[2]) == (10, 10)
+
+  x = rand(1, 10)
+  layer = PairwiseFusion(.+,  Dense(1, 10),  Dense(10, 1))
+  y = layer(x)
+  @test length(y) == 2
+  @test size(y[1]) == (10, 10)
+  @test size(y[2]) == (1, 10)
+end