Gradient of jacobian in Flux v0.10 (zygote)

General Usage
1

Hi all,

in our project on modelling probability distributions we need to calculate a gradient of a function, which involves jacobian of the model. Zygote (Flux) crashes with the output that mutating arrays are not support. Does anyone know, how to do this? Our minimum working example would look like this


function jacobian(f, x)
    m = length(x)
    bf = Zygote.Buffer(x,m, m)
    for i in 1:m
        bf[i, :] = gradient(x -> f(x)[i], x)[1]
    end
    copy(bf)
end

gs = gradient(() -> sum(jacobian(m, x0)), ps)

Thanks for help in advance.
Tomas