Zygote: Mutating arrays is not supported

General Usage
1

I am trying to calculate the laplacian of a neural network using Flux library. My code is:

using Flux
using Zygote

model = Chain(Dense(2,5,sigmoid), Dense(5,1))

function laplace(x)
    a, b = size(x)
    Δ = Zygote.Buffer(zeros(b))
    deriv2 = sum(Diagonal(ones(a*b)).*Zygote.hessian(v -> sum(model(v)), x), dims=1)
    for i=1:b
        for j=1:a
            Δ[i] += deriv2[(i-1)*a+j]
        end
    end
    return copy(Δ)
end

gradient(x -> sum(laplace(x)), rand(2,5))
2

it is indeed not supported

3

I am using Zygote.Buffer.

4

I can’t see at first glance what is causing the error, but note that in any case differentiating loops like this tends to be slow in Zygote. I would drop the Buffer and write

Δ = vec(sum(reshape(deriv2, a, b); dims = 1))

Also, the expression for deriv2 seems to be equivalent to

deriv2 = diag(Zygote.hessian(v -> sum(model(v)), x))

and there may be more efficient ways to compute the diagonal of the hessian with Zygote.