NewRecur experimental interface

src/Fluxperimental.jl

@@ -13,4 +13,6 @@ include("chain.jl")
 
 include("compact.jl")
 
+include("new_recur.jl")
+
 end # module Fluxperimental

src/new_recur.jl

@@ -0,0 +1,105 @@
+
+
+
+##### Helper scan funtion which can likely be put into NNLib. #####
+"""
+  scan
+
+Recreating jax.lax.scan functionality in julia.
+"""
+function scan_full(func, init_carry, xs::AbstractVector{<:AbstractArray})
+  # xs = Flux.eachlastdim(x_block)
+  x_init, x_rest = Iterators.peel(xs)
+
+  (carry, out_) = func(init_carry, x_init)
+
+  init = (typeof(out_)[out_], carry)
+
+  function recurrence_op(input, (outputs, carry))
+    carry, out = func(carry, input)
+    return vcat(outputs, typeof(out)[out]), carry
+  end
+  results = foldr(recurrence_op, xs[(begin+1):end]; init)
+  results[2], results[1]
+end
+
+function scan_full(func, init_carry, x_block)
+  xs_ = Flux.eachlastdim(x_block)
+  xs = if xs_ isa Base.Generator
+    collect(xs_) # eachlastdim produces a generator in non-gradient environment
+  else
+    xs_
+  end
+  scan_full(func, init_carry, xs)
+end
+
+function scan_partial(func, init_carry, xs::AbstractVector{<:AbstractArray})
+  x_init, x_rest = Iterators.peel(xs)
+  (carry, y) = func(init_carry, x_init)
+  for x in x_rest
+    (carry, y) = func(carry, x)
+  end
+  # carry, y
+  carry, y
+end
+
+function scan_partial(func, init_carry, x_block)
+  xs_ = Flux.eachlastdim(x_block)
+  xs = if xs_ isa Base.Generator
+    collect(xs_) # eachlastdim produces a generator in non-gradient environment
+  else
+    xs_
+  end
+  scan_partial(func, init_carry, xs)
+end
+
+
+"""
+  NewRecur
+New Recur. An experimental recur interface for removing statefullness in recurrent architectures for flux.
+"""
+struct NewRecur{RET_SEQUENCE, T}
+  cell::T
+  # state::S
+  function NewRecur(cell; return_sequence::Bool=false)
+    new{return_sequence, typeof(cell)}(cell)
+  end
+  function NewRecur{true}(cell)
+    new{true, typeof(cell)}(cell)
+  end
+  function NewRecur{false}(cell)
+    new{false, typeof(cell)}(cell)
+  end
+end
+
+Flux.@functor NewRecur
+Flux.trainable(a::NewRecur) = (; cell = a.cell)
+Base.show(io::IO, m::NewRecur) = print(io, "Recur(", m.cell, ")")
+NewRNN(a...; return_sequence::Bool=false, ka...) = NewRecur(Flux.RNNCell(a...; ka...); return_sequence=return_sequence)
+
+
+(l::NewRecur)(init_carry, x_mat::AbstractMatrix) = MethodError("Matrix is ambiguous with NewRecur")
+(l::NewRecur)(init_carry, x_mat::AbstractVector{T}) where {T<:Number} = MethodError("Vector is ambiguous with NewRecur")
+
+(l::NewRecur)(xs) = l(l.cell.state0, xs)
+
+
+function (l::NewRecur{false})(init_carry,
+                              xs)
+  results = scan_partial(l.cell, init_carry, xs)
+  results[2]
+end
+
+# From Lux.jl: https://github.com/LuxDL/Lux.jl/pull/287/
+function (l::NewRecur{true})(init_carry,
+                             xs,)
+
+  results = scan_full(l.cell, init_carry, xs)
+
+  h = results[2]
+  sze = size(h[1])
+  reshape(reduce(hcat, h), sze[1], sze[2], length(h))
+end
+
+
+

test/new_recur.jl

@@ -0,0 +1,182 @@
+
+
+@testset "NewRecur RNN" begin
+  @testset "Forward Pass" begin
+    cell = Flux.RNNCell(1, 1, identity)
+    layer = Fluxperimental.NewRecur(cell; return_sequence=true)
+    layer.cell.Wi .= 5.0
+    layer.cell.Wh .= 4.0
+    layer.cell.b .= 0.0f0
+    layer.cell.state0 .= 7.0
+    x = reshape([2.0f0, 3.0f0], 1, 1, 2)
+
+    # @show layer(x)
+    @test eltype(layer(x)) <: Float32
+    @test size(layer(x)) == (1, 1, 2)
+
+    @test_throws MethodError layer([2.0f0])
+    @test_throws MethodError layer([2.0f0;; 3.0f0])
+  end
+
+
+  @testset "gradients-implicit" begin
+    cell = Flux.RNNCell(1, 1, identity)
+    layer = Flux.Recur(cell)
+    layer.cell.Wi .= 5.0
+    layer.cell.Wh .= 4.0
+    layer.cell.b .= 0.0f0
+    layer.cell.state0 .= 7.0
+    x = [[2.0f0], [3.0f0]]
+
+    # theoretical primal gradients
+    primal =
+      layer.cell.Wh .* (layer.cell.Wh * layer.cell.state0 .+ x[1] .* layer.cell.Wi) .+
+      x[2] .* layer.cell.Wi
+    ∇Wi = x[1] .* layer.cell.Wh .+ x[2]
+    ∇Wh = 2 .* layer.cell.Wh .* layer.cell.state0 .+ x[1] .* layer.cell.Wi
+    ∇b = layer.cell.Wh .+ 1
+    ∇state0 = layer.cell.Wh .^ 2
+
+    nm_layer = Fluxperimental.NewRecur(cell; return_sequence = true)
+    ps = Flux.params(nm_layer)
+    x_block = reshape(vcat(x...), 1, 1, length(x))
+    e, g = Flux.withgradient(ps) do
+      out = nm_layer(x_block)
+      sum(out[1, 1, 2])
+    end
+
+    @test primal[1] ≈ e
+    @test ∇Wi ≈ g[ps[1]]
+    @test ∇Wh ≈ g[ps[2]]
+    @test ∇b ≈ g[ps[3]]
+    @test ∇state0 ≈ g[ps[4]]
+  end
+
+
+  @testset "gradients-explicit" begin
+
+    cell = Flux.RNNCell(1, 1, identity)
+    layer = Flux.Recur(cell)
+    layer.cell.Wi .= 5.0
+    layer.cell.Wh .= 4.0
+    layer.cell.b .= 0.0f0
+    layer.cell.state0 .= 7.0
+    x = [[2.0f0], [3.0f0]]
+
+    # theoretical primal gradients
+    primal =
+      layer.cell.Wh .* (layer.cell.Wh * layer.cell.state0 .+ x[1] .* layer.cell.Wi) .+
+      x[2] .* layer.cell.Wi
+    ∇Wi = x[1] .* layer.cell.Wh .+ x[2]
+    ∇Wh = 2 .* layer.cell.Wh .* layer.cell.state0 .+ x[1] .* layer.cell.Wi
+    ∇b = layer.cell.Wh .+ 1
+    ∇state0 = layer.cell.Wh .^ 2
+
+
+    x_block = reshape(vcat(x...), 1, 1, length(x))
+    nm_layer = Fluxperimental.NewRecur(cell; return_sequence = true)
+    e, g = Flux.withgradient(nm_layer) do layer
+      out = layer(x_block)
+      sum(out[1, 1, 2])
+    end
+    grads = g[1][:cell]
+
+    @test primal[1] ≈ e
+    @test ∇Wi ≈ grads[:Wi]
+    @test ∇Wh ≈ grads[:Wh]
+    @test ∇b ≈ grads[:b]
+    @test ∇state0 ≈ grads[:state0]
+
+  end
+end
+
+@testset "New Recur RNN Partial Sequence" begin
+
+  @testset "Forward Pass" begin
+    cell = Flux.RNNCell(1, 1, identity)
+    layer = Fluxperimental.NewRecur(cell)
+    layer.cell.Wi .= 5.0
+    layer.cell.Wh .= 4.0
+    layer.cell.b .= 0.0f0
+    layer.cell.state0 .= 7.0
+    x = reshape([2.0f0, 3.0f0], 1, 1, 2)
+
+    @test eltype(layer(x)) <: Float32
+    @test size(layer(x)) == (1, 1)
+
+    @test_throws MethodError layer([2.0f0])
+    @test_throws MethodError layer([2.0f0;; 3.0f0])
+
+  end
+
+  @testset "gradients-implicit" begin
+    cell = Flux.RNNCell(1, 1, identity)
+    layer = Flux.Recur(cell)
+    layer.cell.Wi .= 5.0
+    layer.cell.Wh .= 4.0
+    layer.cell.b .= 0.0f0
+    layer.cell.state0 .= 7.0
+    x = [[2.0f0], [3.0f0]]
+
+    # theoretical primal gradients
+    primal =
+      layer.cell.Wh .* (layer.cell.Wh * layer.cell.state0 .+ x[1] .* layer.cell.Wi) .+
+      x[2] .* layer.cell.Wi
+    ∇Wi = x[1] .* layer.cell.Wh .+ x[2]
+    ∇Wh = 2 .* layer.cell.Wh .* layer.cell.state0 .+ x[1] .* layer.cell.Wi
+    ∇b = layer.cell.Wh .+ 1
+    ∇state0 = layer.cell.Wh .^ 2
+
+    nm_layer = Fluxperimental.NewRecur(cell; return_sequence = false)
+    ps = Flux.params(nm_layer)
+    x_block = reshape(vcat(x...), 1, 1, length(x))
+    e, g = Flux.withgradient(ps) do
+      out = (nm_layer)(x_block)
+      sum(out)
+    end
+
+    @test primal[1] ≈ e
+    @test ∇Wi ≈ g[ps[1]]
+    @test ∇Wh ≈ g[ps[2]]
+    @test ∇b ≈ g[ps[3]]
+    @test ∇state0 ≈ g[ps[4]]
+  end
+
+  @testset "gradients-explicit" begin
+
+
+    cell = Flux.RNNCell(1, 1, identity)
+    layer = Flux.Recur(cell)
+    layer.cell.Wi .= 5.0
+    layer.cell.Wh .= 4.0
+    layer.cell.b .= 0.0f0
+    layer.cell.state0 .= 7.0
+    x = [[2.0f0], [3.0f0]]
+
+    # theoretical primal gradients
+    primal =
+      layer.cell.Wh .* (layer.cell.Wh * layer.cell.state0 .+ x[1] .* layer.cell.Wi) .+
+      x[2] .* layer.cell.Wi
+    ∇Wi = x[1] .* layer.cell.Wh .+ x[2]
+    ∇Wh = 2 .* layer.cell.Wh .* layer.cell.state0 .+ x[1] .* layer.cell.Wi
+    ∇b = layer.cell.Wh .+ 1
+    ∇state0 = layer.cell.Wh .^ 2
+
+
+    x_block = reshape(vcat(x...), 1, 1, length(x))
+    nm_layer = Fluxperimental.NewRecur(cell; return_sequence = false)
+    e, g = Flux.withgradient(nm_layer) do layer
+      out = layer(x_block)
+      sum(out)
+    end
+    grads = g[1][:cell]
+
+    @test primal[1] ≈ e
+    @test ∇Wi ≈ grads[:Wi]
+    @test ∇Wh ≈ grads[:Wh]
+    @test ∇b ≈ grads[:b]
+    @test ∇state0 ≈ grads[:state0]
+
+  end
+end
+

test/runtests.jl

@@ -8,4 +8,6 @@ using Flux, Fluxperimental
 
   include("compact.jl")
 
+  include("new_recur.jl")
+
 end