full fix to #1352

src/lib/buffer.jl

@@ -1,13 +1,14 @@
 grad_mut(cx::Context, b::Buffer, ::Type=Union{}) =
  _get!(() -> fill!(similar(b.data, Any), nothing), cache(cx), b)
+# S is the eltype we are about to set into the buffer accumulator, so allocte wide enough
 grad_mut(cx::Context, b::Buffer{T}, ::Type{S}=Union{}) where {T<:Number, S<:Number} =
  _get!(() -> fill!(similar(b.data, float(promote_type(T, S))), 0), cache(cx), b)
 
 @non_differentiable Buffer(::Any...)
 
 @adjoint function getindex(b::Buffer, i...)
- b[i...], function (Δ::S) where {S}
- grad = grad_mut(__context__, b, S)
+ b[i...], function (Δ)
+ grad = grad_mut(__context__, b, eltype(Δ))
  grad[i...] = accum(grad[i...], Δ)
  return
  end

test/gradcheck.jl

@@ -1553,46 +1553,13 @@ using Zygote: Buffer
  @test ∇W2 == W2
  @test ∇x == 6 .* x
 
- @testset "incorrect promotion (#1352)" begin
- u = [0.75, 0.5]
- p = [-1.5, 0.05, 0.2, 0.01]
-
- # in-place
- function g1352!(du, u, p, t)
- du[1, 1] = p[3] * u[1] + p[4] * u[2]
- du[1, 2] = p[3] * u[1] + p[4] * u[2]
- du[2, 1] = p[4] * u[1] + p[3] * u[2]
- du[2, 2] = p[4] * u[1] + p[3] * u[2]
- return nothing
- end
- du1_inplace, back_inplace = Zygote.pullback(u, p) do u, p
- du = Zygote.Buffer(Matrix{Float64}(undef, 2, 2))
- g1352!(du, u, p, 1.0)
- return copy(du[:, 1])
- end
-
- # out-of-place
- function g1352(u, p, t)
- du11 = p[3] * u[1] + p[4] * u[2]
- du12 = p[3] * u[1] + p[4] * u[2]
- du21 = p[4] * u[1] + p[3] * u[2]
- du22 = p[4] * u[1] + p[3] * u[2]
- return [du11 du12
- du21 du22]
- end
- du1, back = Zygote.pullback(u, p) do u, p
- du = g1352(u, p, 1.0)
- return du[:, 1]
- end
+ # reduced mwe of #1352
+ @test Zygote.gradient([0,0]) do x
+ buf = Zygote.Buffer(similar(x))
+ buf[:] = x
+ sum(copy(buf[1:2]))
+ end == ([1,1],)
 
- # comparison
- @test du1_inplace ≈ du1
- v = randn(2)
- ∇u_inplace, ∇p_inplace = back_inplace(v)
- ∇u, ∇p = back(v)
- @test ∇u_inplace ≈ ∇u
- @test ∇p_inplace ≈ ∇p
- end
 end
 
 @testset "AbstractArray Addition / Subtraction / Negation" begin