gradients of channels

FluxML · Aug 20, 2019 · e2fcd61 · e2fcd61 · MikeInnes · Aug 20, 2019
1 parent 1a049a2
commit e2fcd61
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 1 deletion.
diff --git a/src/compiler/reverse.jl b/src/compiler/reverse.jl
@@ -84,7 +84,7 @@ function instrument(ir::IR)
   pr = Pipe(ir)
   for (v, st) in pr
     ex = st.expr
-    isexpr(ex, :foreigncall) && continue
+    isexpr(ex, :foreigncall, :isdefined) && continue
     isexpr(ex, :enter, :leave) && error("try/catch is not supported.")
     ex = instrument_new!(pr, v, ex)
     ex = instrument_literals!(pr, v, ex)

diff --git a/src/lib/base.jl b/src/lib/base.jl
@@ -46,3 +46,23 @@ end
     (nothing, Δ, nothing)
   end
 end
+
+# Channels
+
+@nograd Channel
+
+grad_mut(ch::Channel) = Channel(ch.sz_max)
+
+@adjoint! function put!(ch::Channel, x)
+  put!(ch, x), function (ȳ)
+    x̄ = take!(grad_mut(__context__, ch))
+    (nothing, accum(x̄, ȳ), nothing)
+  end
+end
+
+@adjoint! function take!(ch::Channel)
+  take!(ch), function (x̄)
+    put!(grad_mut(__context__, ch), x̄)
+    return
+  end
+end
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -13,6 +13,10 @@ end
   include("features.jl")
 end
 
+@testset "Data Structures" begin
+  include("structures.jl")
+end
+
 @testset "Gradients" begin
   include("gradcheck.jl")
 end

diff --git a/test/structures.jl b/test/structures.jl
@@ -0,0 +1,8 @@
+using Zygote, Test
+
+function f(x)
+  ch = Channel(Inf)
+  put!(ch, x^2) + take!(ch)
+end
+
+@test gradient(f, 5) == (20,)