Zygote: Mutating arrays is not supported

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]
    return copy(Δ)

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

it is indeed not supported

I am using Zygote.Buffer.

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.