Merge pull request #360 from FluxML/buffer-broadcast

In-place broadcast over buffers

src/lib/buffer.jl

@@ -51,6 +51,11 @@ function Base.setindex!(b::Buffer, v, i...)
   b.data[i...] = v
 end
 
+function Base.copyto!(b::Buffer, data)
+  b.freeze && error("Buffer is frozen")
+  copyto!(b.data, data)
+end
+
 function Base.copy(b::Buffer)
   b.freeze = true
   return b.data
@@ -85,6 +90,18 @@ end
   end
 end
 
+@adjoint! function copyto!(b::Buffer, xs)
+  copyto!(b, xs), function (_)
+    grad = grad_mut(__context__, b)
+    x̄s = copy(grad)
+    grad .= eltype(grad) <: Number ? 0 : nothing
+    return (nothing, x̄s)
+  end
+end
+
+_pullback(cx::AContext, ::typeof(Broadcast.materialize!), b::Buffer, x::AbstractArray) =
+  _pullback(cx, copyto!, b, x)
+
 @adjoint function copy(b::Buffer)
   copy(b), function (b̄)
     grad_mut(__context__, b)[:] = b̄

test/gradcheck.jl

@@ -840,6 +840,12 @@ using Zygote: Buffer
   @test eachindex(buf) == 1:3
   @test stride(buf, 2) === 3
   @test strides(buf) === (1, )
+
+  @test gradient([1, 2, 3]) do xs
+    b = Zygote.Buffer(xs)
+    b .= xs .* 2
+    return sum(copy(b))
+  end == ([2,2,2],)
 end
 
 @testset "FillArrays" begin