# Autograd through loss function with derivatives in

**URL:** <https://discourse.julialang.org/t/autograd-through-loss-function-with-derivatives-in/57864>\
**Category:** Machine Learning\
**Created:** [March 24, 2021, 12:09pm UTC](https://discourse.julialang.org/t/autograd-through-loss-function-with-derivatives-in/57864 "2021-03-24T12:09:45Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![RynoLaubscher](https://avatars.discourse-cdn.com/v4/letter/r/8baadc/32.png) [@RynoLaubscher](https://discourse.julialang.org/u/RynoLaubscher)\
**Post date:** [March 24, 2021, 12:09pm UTC](https://discourse.julialang.org/t/autograd-through-loss-function-with-derivatives-in/57864/1 "2021-03-24T12:09:45Z")

</div>

Hi All, I am relatively new to Julia, coming from a C/C++ & Python background. I am absolutely loving Julia, but I am struggling a bit to build a simple PINN model using Zygote/Flux (and I tried ForwardDiff). I am using one of the standard examples from Raissi et al. [Simple PINN NLS](https://github.com/maziarraissi/PINNs), but I simplified the loss function for testing purposes. I would like to define a loss function which has derivatives with respect to input features, and then use Zygote.gradient to find the gradients of the loss function with respect to the model parameters.  
Please see below the definition of the network:

```julia
#define layers
fc1 = Dense(2, 20, tanh);
fc2 = Dense(20, 2, identity);

#create chain and set optimizer
model = Chain(fc1, fc2);
pars = Flux.params(model);
opt = ADAM();

```

Here is my loss function code:

```julia
function loss_fn(model, X)
    u(x) = model(x)[1,:]
    v(x) = model(x)[2,:]
    ∂u∂x = Zygote.gradient(a -> sum(u(a)), X)[1][1,:]
    ∂u∂t = Zygote.gradient(a -> sum(u(a)), X)[1][2,:]
    ∂v∂x = Zygote.gradient(a -> sum(v(a)), X)[1][1,:]
    ∂v∂t = Zygote.gradient(a -> sum(v(a)), X)[1][2,:]
    loss = mean(u(X).^2) + mean(v(X).^2) + mean(∂u∂x.^2) +
            mean(∂u∂t.^2) + mean(∂v∂x.^2) + mean(∂v∂t.^2)
    return loss
end

```

Here is the line I am trying to use to find the gradients of the loss function to model parameters:

```julia
grads = Flux.gradient(() -> loss_fn(model, x0_t0), pars)

```

But when I run the above code I get a “mutating” array error. I know other colleagues have proposed using a combination of ReverseDiff and Zygote.gradient to fix this. But honestly I am not sure what that would look like? Any help would be appreciated.

---

<div class="post-metadata">

**Author:** ![RynoLaubscher](https://avatars.discourse-cdn.com/v4/letter/r/8baadc/32.png) [@RynoLaubscher](https://discourse.julialang.org/u/RynoLaubscher)\
**Post date:** [March 25, 2021, 6:51am UTC](https://discourse.julialang.org/t/autograd-through-loss-function-with-derivatives-in/57864/2 "2021-03-25T06:51:46Z")

</div>

Looks like using ReverseDiff over Zygote works, see this discussion [post](https://discourse.julialang.org/t/how-to-use-gradient-of-neural-network-as-the-loss-function/50569/14) for the discussion. Here is my loss function with using Zygote to get the gradients:

```julia
layer(W, b, x) = tanh.(W * x .+ b);
function loss(W, b, x)
    u(x) = layer(W, b, x)[1,:]
    v(x) = layer(W, b, x)[2,:]
    ∂u∂x = Zygote.gradient(a -> sum(u(a)), x)[1][1,:]
    ∂v∂x = Zygote.gradient(a -> sum(v(a)), x)[1][1,:]
    mean(∂u∂x.^2) + mean(∂v∂x.^2) #just a simple loss function, not physical
end

```

Here is the gradient of the loss function with respect to the model weights and biases:

```julia
const W, b = rand(2,2), rand(2);
const input = (W, b);
const diff = ReverseDiff.compile(ReverseDiff.GradientTape((a,b) -> loss(a,b,x0_t0), input));
a = ReverseDiff.gradient!(map(copy, input), diff, input);
a[1]
2×2 Array{Float64,2}:
 0.118311 0.0
 0.126966 0.0

```