# Errors when learning a function of a Jacobian. What am I doing wrong?

**URL:** https://discourse.julialang.org/t/errors-when-learning-a-function-of-a-jacobian-what-am-i-doing-wrong/76352
**Category:** Machine Learning
**Tags:** flux, zygote
**Created:** [February 13, 2022, 2:23pm UTC](https://discourse.julialang.org/t/errors-when-learning-a-function-of-a-jacobian-what-am-i-doing-wrong/76352 "2022-02-13T14:23:20Z")
**Posts on this page:** 1
**Page:** 1

<div class="post-metadata">

### Author: ![ynsch](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ynsch/32/29558_2.png) [@ynsch](https://discourse.julialang.org/u/ynsch)
#### Post date: [February 13, 2022, 2:23pm UTC](https://discourse.julialang.org/t/errors-when-learning-a-function-of-a-jacobian-what-am-i-doing-wrong/76352/1 "2022-02-13T14:23:20Z")

</div>

Hello,

I am working on a simple example with Flux and Zygote, mainly translating this [notebook](https://github.com/acids-ircam/diffusion_models/blob/main/diffusion_01_score.ipynb) to Julia.

The loss function contains the divergence (trace of the Jacobian) and Zygote gives an error when updating the model.

```julia
import ManifoldLearning: swiss_roll
using Flux
#using CUDA
using Zygote
using LinearAlgebra

# 2x1000 data
data = swiss_roll(1000,1.0)[1][[1,3],:] #|> gpu 

model = Chain( Dense(2, 128, relu),
                Dense(128,128,relu),
                  Dense(128, 2,relu)) #|> gpu

opt = ADAM(1e-3)

function score_matching(model,data)
    logp = model(data)
    norm_loss = sum(abs2,logp)/2
    div = tr(jacobian(model,data)[1])
    return 0.5(norm_loss+div)
end

ps = Flux.params(model)

function train_step!(model,data,ps,opt)
    gs = gradient(() -> score_matching(model, data), ps)
    Flux.Optimise.update!(opt, ps, gs)
end

train_step!(model,data,ps,opt)

```

gives

```julia
LoadError: Mutating arrays is not supported -- called copyto!(::SubArray{Float64, 1, Matrix{Float64}, Tuple{Int64, Base.Slice{Base.OneTo{Int64}}}, true}, _...)

```

The stacktrace leads to `gradcopy` in `Zygote.withjacobian`.  
Is it not possible to differentiate a function of the Jacobian?

Even more surprising, when I use `CUDA` (by including comments), I get a different error:

```julia
LoadError: Compiling Tuple{CUDA.var"##context!#59", Bool, typeof(context!), CUDA.var"#216#217"{CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, UInt32, DataType}, CuContext}: try/catch is not supported.

```

The try/catch seems to come from `fill!` used by `_eyelike` in `Zygote.withjacobian`.  
Should I avoid using the Jacobian somehow? I also tried using `ForwardDiff.jacobian` instead, but I get yet another error.