# How to use gradient of neural network as the loss function?

**URL:** <https://discourse.julialang.org/t/how-to-use-gradient-of-neural-network-as-the-loss-function/50569>\
**Category:** Machine Learning\
**Tags:** question\
**Created:** [November 22, 2020, 6:19am UTC](https://discourse.julialang.org/t/how-to-use-gradient-of-neural-network-as-the-loss-function/50569 "2020-11-22T06:19:50Z")\
**Posts on this page:** 14\
**Page:** 1

<div class="post-metadata">

**Author:** ![XiaodongMa-MRI](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/xiaodongma-mri/32/16171_2.png) [@XiaodongMa-MRI](https://discourse.julialang.org/u/XiaodongMa-MRI)\
**Post date:** [November 22, 2020, 6:19am UTC](https://discourse.julialang.org/t/how-to-use-gradient-of-neural-network-as-the-loss-function/50569/1 "2020-11-22T06:19:50Z")

</div>

Dear all,

I want to put the gradient of neural network (i.e., the gradient of NN output w.r.t input) in the loss function. However, it seems that Zygote cannot differentiate that. Here is my code:

```julia
using Flux
using Zygote
using Random

Xtrain = rand(Float32,(4,100));
Ytrain = rand(Float32,(1,100));

# *******************NN********************
model = Chain(x->x[1:2,:],
	      Dense(2, 50, tanh),
              Dense(50, 50, tanh),
              Dense(50, 1))
θ = params(model);

# *******************Lost function********************
function loss(x,y)
    ps = Flux.Params([x]);
    g = Flux.gradient(ps) do
        model(x)[1]
    end
    return ( Flux.mse(g[x][1:2,:][1] , x[3:4,:]) + Flux.mse(model(x),y) )
end

# *******************Training********************
train_loader = Flux.Data.DataLoader(Xtrain,Ytrain, batchsize=1,shuffle=true);

Nepochs = 10;
opt = ADAM(0.001)

for epoch in 1:Nepochs
  @time Flux.train!(loss, θ, train_loader, opt)
end

```

However, the code failed to work. This is error message:

```julia
ERROR: Can't differentiate foreigncall expression
Stacktrace:
 [1] error(::String) at ./error.jl:33
 [2] get at ./iddict.jl:87 [inlined]
 [3] (::typeof(∂(get)))(::Nothing) at /home/wuxxx184/ma000311/.julia/packages/Zygote/chgvX/src/compiler/interface2.jl:0
 [4] accum_global at /home/wuxxx184/ma000311/.julia/packages/Zygote/chgvX/src/lib/lib.jl:56 [inlined]
 [5] (::typeof(∂(accum_global)))(::Nothing) at /home/wuxxx184/ma000311/.julia/packages/Zygote/chgvX/src/compiler/interface2.jl:0
 [6] #89 at /home/wuxxx184/ma000311/.julia/packages/Zygote/chgvX/src/lib/lib.jl:67 [inlined]
 [7] (::typeof(∂(λ)))(::Nothing) at /home/wuxxx184/ma000311/.julia/packages/Zygote/chgvX/src/compiler/interface2.jl:0
 [8] #1550#back at /home/wuxxx184/ma000311/.julia/packages/ZygoteRules/6nssF/src/adjoint.jl:49 [inlined]
 [9] (::typeof(∂(λ)))(::Nothing) at /home/wuxxx184/ma000311/.julia/packages/Zygote/chgvX/src/compiler/interface2.jl:0
 [10] getindex at ./tuple.jl:24 [inlined]
 [11] gradindex at /home/wuxxx184/ma000311/.julia/packages/Zygote/chgvX/src/compiler/reverse.jl:12 [inlined]
 [12] #3 at ./REPL[8]:5 [inlined]
 [13] (::typeof(∂(λ)))(::Nothing) at /home/wuxxx184/ma000311/.julia/packages/Zygote/chgvX/src/compiler/interface2.jl:0
 [14] #54 at /home/wuxxx184/ma000311/.julia/packages/Zygote/chgvX/src/compiler/interface.jl:177 [inlined]
 [15] (::typeof(∂(λ)))(::Nothing) at /home/wuxxx184/ma000311/.julia/packages/Zygote/chgvX/src/compiler/interface2.jl:0
 [16] gradient at /home/wuxxx184/ma000311/.julia/packages/Zygote/chgvX/src/compiler/interface.jl:54 [inlined]
 [17] (::typeof(∂(gradient)))(::Nothing) at /home/wuxxx184/ma000311/.julia/packages/Zygote/chgvX/src/compiler/interface2.jl:0
 [18] loss at ./REPL[8]:4 [inlined]
 [19] (::typeof(∂(loss)))(::Float32) at /home/wuxxx184/ma000311/.julia/packages/Zygote/chgvX/src/compiler/interface2.jl:0
 [20] #145 at /home/wuxxx184/ma000311/.julia/packages/Zygote/chgvX/src/lib/lib.jl:175 [inlined]
 [21] #1681#back at /home/wuxxx184/ma000311/.julia/packages/ZygoteRules/6nssF/src/adjoint.jl:49 [inlined]
 [22] #15 at /home/wuxxx184/ma000311/.julia/packages/Flux/05b38/src/optimise/train.jl:83 [inlined]
 [23] (::typeof(∂(λ)))(::Float32) at /home/wuxxx184/ma000311/.julia/packages/Zygote/chgvX/src/compiler/interface2.jl:0
 [24] (::Zygote.var"#54#55"{Params,Zygote.Context,typeof(∂(λ))})(::Float32) at /home/wuxxx184/ma000311/.julia/packages/Zygote/chgvX/src/compiler/interface.jl:177
 [25] gradient(::Function, ::Params) at /home/wuxxx184/ma000311/.julia/packages/Zygote/chgvX/src/compiler/interface.jl:54
 [26] macro expansion at /home/wuxxx184/ma000311/.julia/packages/Flux/05b38/src/optimise/train.jl:82 [inlined]
 [27] macro expansion at /home/wuxxx184/ma000311/.julia/packages/Juno/n6wyj/src/progress.jl:134 [inlined]
 [28] train!(::Function, ::Params, ::Flux.Data.DataLoader{Tuple{Array{Float32,2},Array{Float32,2}}}, ::ADAM; cb::Flux.Optimise.var"#16#22") at /home/wuxxx184/ma000311/.julia/packages/Flux/05b38/src/optimise/train.jl:80
 [29] train!(::Function, ::Params, ::Flux.Data.DataLoader{Tuple{Array{Float32,2},Array{Float32,2}}}, ::ADAM) at /home/wuxxx184/ma000311/.julia/packages/Flux/05b38/src/optimise/train.jl:78
 [30] macro expansion at ./timing.jl:174 [inlined]
 [31] top-level scope at ./REPL[12]:2

```

Does anyone know how to fix this?

Thanks a lot!  
Xiaodong

---

<div class="post-metadata">

**Author:** ![Tomas\_Pevny](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/tomas_pevny/32/25466_2.png) [@Tomas\_Pevny](https://discourse.julialang.org/u/Tomas_Pevny)\
**Post date:** [November 22, 2020, 9:04am UTC](https://discourse.julialang.org/t/how-to-use-gradient-of-neural-network-as-the-loss-function/50569/2 "2020-11-22T09:04:17Z")

</div>

I am doing someting similar to you and I have made it work, but it is not something trivial and I would not recommend this to begginers. You need to ensure that all gradients (aka Pullbacks) in your inner gradient are twice differentiable. I would recommend to make the problem example smaller (replace train by a gradient call), etc. I am not sure that the `getindex` is for example twice differentiable. So may-be replacing `model(x)[1]` with `sum(model(x))` might do the trick.

---

<div class="post-metadata">

**Author:** ![XiaodongMa-MRI](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/xiaodongma-mri/32/16171_2.png) [@XiaodongMa-MRI](https://discourse.julialang.org/u/XiaodongMa-MRI)\
**Post date:** [November 22, 2020, 3:51pm UTC](https://discourse.julialang.org/t/how-to-use-gradient-of-neural-network-as-the-loss-function/50569/3 "2020-11-22T15:51:56Z")

</div>

Thank you very much for your suggesions, Tomas. I tried to remove all the indexes in the loss function:

```julia
function loss(x,y)
    ps = Flux.Params([x]);
    g = Flux.gradient(ps) do
        sum(model(x))
    end
    return ( sum(Flux.mse(g[x])) )
end

```

But same error will occur.  
I will replace train by a simpler gradient call to see if I can trace Pullbacks.

Xiaodong

---

<div class="post-metadata">

**Author:** ![Tomas\_Pevny](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/tomas_pevny/32/25466_2.png) [@Tomas\_Pevny](https://discourse.julialang.org/u/Tomas_Pevny)\
**Post date:** [November 22, 2020, 4:19pm UTC](https://discourse.julialang.org/t/how-to-use-gradient-of-neural-network-as-the-loss-function/50569/4 "2020-11-22T16:19:07Z")

</div>

If I will not forget, I will send you tomorrow a hacked ZygoteRules which trace which pullback is called. It is a bit low-level, but it help me to debig nans and all these things.

---

<div class="post-metadata">

**Author:** ![mkalia94](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mkalia94/32/22849_2.png) [@mkalia94](https://discourse.julialang.org/u/mkalia94)\
**Post date:** [March 12, 2021, 1:16pm UTC](https://discourse.julialang.org/t/how-to-use-gradient-of-neural-network-as-the-loss-function/50569/5 "2021-03-12T13:16:24Z")

</div>

Hi!  
I have the same issue, to no avail. @Tomas_Pevny could you share your solution here?

---

<div class="post-metadata">

**Author:** ![Tomas\_Pevny](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/tomas_pevny/32/25466_2.png) [@Tomas\_Pevny](https://discourse.julialang.org/u/Tomas_Pevny)\
**Post date:** [March 15, 2021, 8:57am UTC](https://discourse.julialang.org/t/how-to-use-gradient-of-neural-network-as-the-loss-function/50569/6 "2021-03-15T08:57:06Z")

</div>

Can you post MWE?

---

<div class="post-metadata">

**Author:** ![mkalia94](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mkalia94/32/22849_2.png) [@mkalia94](https://discourse.julialang.org/u/mkalia94)\
**Post date:** [March 15, 2021, 10:03am UTC](https://discourse.julialang.org/t/how-to-use-gradient-of-neural-network-as-the-loss-function/50569/7 "2021-03-15T10:03:51Z")

</div>

Here’s what I’ve been trying to do:

```julia
using Zygote, Flux
NN = Chain(Conv((3, 3), 1=>1, pad=(1,1), relu), flatten, Dense(784,2)) |> gpu
loss(x) = sum(abs2,Zygote.forward_jacobian(NN,x)[2]) # the second element gives the Jacobian matrix
x = rand(28,28,1,1) |> gpu
loss(x) # works
ps = Flux.params(NN)
res,back = Zygote.pullback(ps) do
 loss(x)
end

```

This gives me a `CuArray only supports bits types` error.

Replacing `forward_jacobian` by `Zygote.jacobian` (this is a new function present in the `Zygote` master branch) in `loss` by doing

```julia
loss(x) = sum(abs2,Zygote.jacobian(NN,x)[1])

```

gives the following error upon training:

```julia
ERROR: this intrinsic must be compiled to be called                                                                                                                                                         
Stacktrace:                                                                                                                                                                                                 
 [1] macro expansion at /home/kaliam/.julia/packages/Zygote/IsBxF/src/compiler/interface2.jl:0 [inlined]                                                                                                    
 [2] _pullback(::Zygote.Context, ::Core.IntrinsicFunction, ::String, ::Type{UInt64}, ::Type{Tuple{Ptr{UInt64}}}, ::Ptr{UInt64}) at /home/kaliam/.julia/packages/Zygote/IsBxF/src/compiler/interface2.jl:9   
 [3] getindex at ./atomics.jl:358 [inlined]                                                                                                                                                                 
 [4] _pullback(::Zygote.Context, ::typeof(getindex), ::Base.Threads.Atomic{UInt64}) at /home/kaliam/.julia/packages/Zygote/IsBxF/src/compiler/interface2.jl:0                                               
 [5] macro expansion at /home/kaliam/.julia/packages/CUDA/wTQsK/lib/utils/call.jl:37 [inlined]                                                                                                              
 [6] cudnnGetVersion at /home/kaliam/.julia/packages/CUDA/wTQsK/lib/cudnn/libcudnn.jl:5 [inlined]                                                                                                           
 [7] _pullback(::Zygote.Context, ::typeof(CUDA.CUDNN.cudnnGetVersion)) at /home/kaliam/.julia/packages/Zygote/IsBxF/src/compiler/interface2.jl:0                                                            
 [8] version at /home/kaliam/.julia/packages/CUDA/wTQsK/lib/cudnn/base.jl:21 [inlined]                                                                                                                      
 [9] _pullback(::Zygote.Context, ::typeof(CUDA.CUDNN.version)) at /home/kaliam/.julia/packages/Zygote/IsBxF/src/compiler/interface2.jl:0                                                               
 [10] #∇conv_filter!#586 at /home/kaliam/.julia/packages/CUDA/wTQsK/lib/cudnn/nnlib.jl:185 [inlined]                                                                                                       
 [11] _pullback(::Zygote.Context, ::CUDA.CUDNN.var"##∇conv_filter!#586", ::Int64, ::Int64, ::typeof(∇conv_filter!), ::CuArray{Float32,4}, ::CuArray{Float32,4}, ::CuArray{Float32,4}, ::DenseConvDims{2,(3, 
3),1,1,(1, 1),(1, 1, 1, 1),(1, 1),false}) at /home/kaliam/.julia/packages/Zygote/IsBxF/src/compiler/interface2.jl:0                                                                                    
 [12] ∇conv_filter! at /home/kaliam/.julia/packages/CUDA/wTQsK/lib/cudnn/nnlib.jl:185 [inlined]                                                                                                            
 [13] _pullback(::Zygote.Context, ::typeof(∇conv_filter!), ::CuArray{Float32,4}, ::CuArray{Float32,4}, ::CuArray{Float32,4}, ::DenseConvDims{2,(3, 3),1,1,(1, 1),(1, 1, 1, 1),(1, 1),false}) at /home/kaliam
/.julia/packages/Zygote/IsBxF/src/compiler/interface2.jl:0                                                                                                                                             
 [14] #∇conv_filter#89 at /home/kaliam/.julia/packages/NNlib/uX1eA/src/conv.jl:116 [inlined]                                                                                                               
 [15] _pullback(::Zygote.Context, ::NNlib.var"##∇conv_filter#89", ::Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}}, ::typeof(∇conv_filter), ::CuArray{Float32,4}, ::CuArray{Float32,4}
, ::DenseConvDims{2,(3, 3),1,1,(1, 1),(1, 1, 1, 1),(1, 1),false}) at /home/kaliam/.julia/packages/Zygote/IsBxF/src/compiler/interface2.jl:0                                                            
 [16] ∇conv_filter at /home/kaliam/.julia/packages/NNlib/uX1eA/src/conv.jl:114 [inlined]                                                                                                                   
 [17] _pullback(::Zygote.Context, ::typeof(∇conv_filter), ::CuArray{Float32,4}, ::CuArray{Float32,4}, ::DenseConvDims{2,(3, 3),1,1,(1, 1),(1, 1, 1, 1),(1, 1),false}) at /home/kaliam/.julia/packages/Zygote
/IsBxF/src/compiler/interface2.jl:0                                                                                                                   
 [18] adjoint at /home/kaliam/.julia/packages/Zygote/IsBxF/src/lib/lib.jl:188 [inlined]                                                                                                            
 [19] _pullback at /home/kaliam/.julia/packages/ZygoteRules/OjfTt/src/adjoint.jl:57 [inlined]                                                                                                          
 [20] #176 at /home/kaliam/.julia/packages/NNlib/uX1eA/src/conv.jl:229 [inlined]                                                                                                                           
 [21] _pullback(::Zygote.Context, ::NNlib.var"#176#179"{Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},CuArray{Float32,4},DenseConvDims{2,(3, 3),1,1,(1, 1),(1, 1, 1, 1),(1, 1),false}
}) at /home/kaliam/.julia/packages/Zygote/IsBxF/src/compiler/interface2.jl:0                                                                                                                           
 [22] Thunk at /home/kaliam/.julia/packages/ChainRulesCore/7d1hl/src/differentials/thunks.jl:98 [inlined]                                                                                                 
 [23] _pullback(::Zygote.Context, ::ChainRulesCore.Thunk{NNlib.var"#176#179"{Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},CuArray{Float32,4},DenseConvDims{2,(3, 3),1,1,(1, 1),(1, 1
, 1, 1),(1, 1),false}}}) at /home/kaliam/.julia/packages/Zygote/IsBxF/src/compiler/interface2.jl:0                                                                                                 
 [24] unthunk at /home/kaliam/.julia/packages/ChainRulesCore/7d1hl/src/differentials/thunks.jl:99 [inlined]                                                                                            
 [25] wrap_chainrules_output at /home/kaliam/.julia/packages/Zygote/IsBxF/src/compiler/chainrules.jl:41 [inlined]                                                                                          
 [26] (::Zygote.var"#534#538"{Zygote.Context,typeof(Zygote.wrap_chainrules_output)})(::ChainRulesCore.Thunk{NNlib.var"#176#179"{Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},CuArray
{Float32,4},DenseConvDims{2,(3, 3),1,1,(1, 1),(1, 1, 1, 1),(1, 1),false}}}) at /home/kaliam/.julia/packages/Zygote/IsBxF/src/lib/array.jl:181                                          
 [27] map at ./tuple.jl:159 [inlined]                                                                                                                                                  
 [28] map at ./tuple.jl:160 [inlined]                                                                                                        
 [29] ∇map at /home/kaliam/.julia/packages/Zygote/IsBxF/src/lib/array.jl:181 [inlined]                                                                                                       
 [30] adjoint at /home/kaliam/.julia/packages/Zygote/IsBxF/src/lib/array.jl:197 [inlined]             
 [31] _pullback at /home/kaliam/.julia/packages/ZygoteRules/OjfTt/src/adjoint.jl:57 [inlined]         
 [32] wrap_chainrules_output at /home/kaliam/.julia/packages/Zygote/IsBxF/src/compiler/chainrules.jl:42 [inlined]                                                                                           
 [33] ZBack at /home/kaliam/.julia/packages/Zygote/IsBxF/src/compiler/chainrules.jl:77 [inlined]      
 [34] _pullback(::Zygote.Context, ::Zygote.ZBack{NNlib.var"#conv_pullback#177"{Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},CuArray{Float32,4},CuArray{Float32,4},DenseConvDims{2,(3
, 3),1,1,(1, 1),(1, 1, 1, 1),(1, 1),false}}}, ::CuArray{Float32,4}) at /home/kaliam/.julia/packages/Zygote/IsBxF/src/compiler/interface2.jl:0                                                               
 [35] Conv at /home/kaliam/.julia/packages/Flux/goUGu/src/layers/conv.jl:147 [inlined]                
 [36] _pullback(::Zygote.Context, ::typeof(∂(λ)), ::CuArray{Float32,4}) at /home/kaliam/.julia/packages/Zygote/IsBxF/src/compiler/interface2.jl:0                                                           
 [37] applychain at /home/kaliam/.julia/packages/Flux/goUGu/src/layers/basic.jl:36 [inlined]          
 [38] _pullback(::Zygote.Context, ::typeof(∂(applychain)), ::CuArray{Float32,2}) at /home/kaliam/.julia/packages/Zygote/IsBxF/src/compiler/interface2.jl:0                                                  
 [39] Chain at /home/kaliam/.julia/packages/Flux/goUGu/src/layers/basic.jl:38 [inlined]               
 [40] _pullback(::Zygote.Context, ::typeof(∂(λ)), ::CuArray{Float32,2}) at /home/kaliam/.julia/packages/Zygote/IsBxF/src/compiler/interface2.jl:0                                                           
 [41] #178 at /home/kaliam/.julia/packages/Zygote/IsBxF/src/lib/lib.jl:191 [inlined]                  
 [42] _pullback(::Zygote.Context, ::Zygote.var"#178#179"{typeof(∂(λ)),Tuple{Tuple{Nothing}}}, ::CuArray{Float32,2}) at /home/kaliam/.julia/packages/Zygote/IsBxF/src/compiler/interface2.jl:0               
 [43] #1698#back at /home/kaliam/.julia/packages/ZygoteRules/OjfTt/src/adjoint.jl:59 [inlined]        
 [44] _pullback(::Zygote.Context, ::Zygote.var"#1698#back#180"{Zygote.var"#178#179"{typeof(∂(λ)),Tuple{Tuple{Nothing}}}}, ::CuArray{Float32,2}) at /home/kaliam/.julia/packages/Zygote/IsBxF/src/compiler/in
terface2.jl:0                                      
 [45] #62 at ./operators.jl:875 [inlined]                                                             
 [46] _pullback(::Zygote.Context, ::typeof(∂(λ)), ::CuArray{Float32,1}) at /home/kaliam/.julia/packages/Zygote/IsBxF/src/compiler/interface2.jl:0                                                           
 [47] #41 at /home/kaliam/.julia/packages/Zygote/IsBxF/src/compiler/interface.jl:41 [inlined]         
 [48] _pullback(::Zygote.Context, ::Zygote.var"#41#42"{typeof(∂(λ))}, ::CuArray{Float32,1}) at /home/kaliam/.julia/packages/Zygote/IsBxF/src/compiler/interface2.jl:0                                       
 [49] jacobian at /home/kaliam/.julia/packages/Zygote/IsBxF/src/lib/grad.jl:148 [inlined]             
 [50] _pullback(::Zygote.Context, ::typeof(jacobian), ::Chain{Tuple{Conv{2,2,typeof(relu),CuArray{Float32,4},CuArray{Float32,1}},typeof(flatten),Dense{typeof(identity),CuArray{Float32,2},CuArray{Float32,1
}}}}, ::CuArray{Float32,4}) at /home/kaliam/.julia/packages/Zygote/IsBxF/src/compiler/interface2.jl:0 
 [51] loss at ./REPL[34]:1 [inlined]               
 [52] _pullback(::Zygote.Context, ::typeof(loss), ::CuArray{Float32,4}) at /home/kaliam/.julia/packages/Zygote/IsBxF/src/compiler/interface2.jl:0                                                           
 [53] #13 at ./REPL[36]:2 [inlined]                
 [54] _pullback(::Zygote.Context, ::var"#13#14") at /home/kaliam/.julia/packages/Zygote/IsBxF/src/compiler/interface2.jl:0                                                                                  
 [55] pullback(::Function, ::Params) at /home/kaliam/.julia/packages/Zygote/IsBxF/src/compiler/interface.jl:247  

```

Thanks for helping out!

---

<div class="post-metadata">

**Author:** ![mkalia94](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mkalia94/32/22849_2.png) [@mkalia94](https://discourse.julialang.org/u/mkalia94)\
**Post date:** [March 15, 2021, 10:05am UTC](https://discourse.julialang.org/t/how-to-use-gradient-of-neural-network-as-the-loss-function/50569/8 "2021-03-15T10:05:28Z")

</div>

Also see [this related Flux issue](https://github.com/FluxML/Flux.jl/issues/1464).

---

<div class="post-metadata">

**Author:** ![arlk](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/arlk/32/2465_2.png) [@arlk](https://discourse.julialang.org/u/arlk)\
**Post date:** [March 18, 2021, 6:58pm UTC](https://discourse.julialang.org/t/how-to-use-gradient-of-neural-network-as-the-loss-function/50569/9 "2021-03-18T18:58:10Z")

</div>

Nested AD is not yet fully supported in Zygote. Your best bet is to set it up the way DiffEqFlux does [it](https://github.com/SciML/DiffEqFlux.jl/blob/master/src/DiffEqFlux.jl#L57) using Zygote over ForwardDiff. For example I need jacobian vector products, i.e. directional derivatives, in my loss function and I set it up this way:

```julia
jvp(f,x,v) = ForwardDiff.partials.(f(ForwardDiff.Dual.(x, v)), 1)

```

You will additionally need to de-dualize the training gradients before you update the neural network weights. One way you can do this:

```julia
import Flux.update!
function update!(opt, x, x̄::AbstractArray{<:ForwardDiff.Dual})
    x̄ = getindex.(ForwardDiff.partials.(x̄), 1)
    Flux.update!(opt, x, x̄)
end

```

---

<div class="post-metadata">

**Author:** ![mkalia94](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mkalia94/32/22849_2.png) [@mkalia94](https://discourse.julialang.org/u/mkalia94)\
**Post date:** [March 23, 2021, 12:31pm UTC](https://discourse.julialang.org/t/how-to-use-gradient-of-neural-network-as-the-loss-function/50569/10 "2021-03-23T12:31:36Z")

</div>

Very cool! This works perfectly for various `Dense` and `Conv` types of neural networks on the CPU, albeit very slow with `ForwardDiff` types in `Conv`. I assume you still need to use `Zygote.pullback` to get the gradients. However, when I use the GPU, I start getting different errors based on the NN I use, which tells me that I need to define more adjoints:

Using `Chain(Dense(10,10))` gives the error:

```julia
ERROR: ArgumentError: invalid index: Dual{Nothing}(1,0,1) of type ForwardDiff.Dual{Nothing,Int64,2}

```

and using `Chain(Dense(10,10,elu))` gives the error:

```julia
ERROR: MethodError: no method matching Float64(::ForwardDiff.Dual{Nothing,Float32,2})

```

Here’s the full code:

```julia
using Zygote, Flux, ZygoteRules, ForwardDiff

# From DiffEqFlux
ZygoteRules.@adjoint function ForwardDiff.Dual{T}(x, ẋ::Tuple) where T
  @assert length(ẋ) == 1
  ForwardDiff.Dual{T}(x, ẋ), ḋ -> (ḋ.partials[1], (ḋ.value,))
end

ZygoteRules.@adjoint ZygoteRules.literal_getproperty(d::ForwardDiff.Dual{T}, ::Val{:partials}) where T =
  d.partials, ṗ -> (ForwardDiff.Dual{T}(ṗ[1], 0),)

ZygoteRules.@adjoint ZygoteRules.literal_getproperty(d::ForwardDiff.Dual{T}, ::Val{:value}) where T =
  d.value, ẋ -> (ForwardDiff.Dual{T}(0, ẋ),)

# Define neural nets, and vecs

f = Chain(Dense(100,10)) |> gpu
x = rand(100) |> gpu
v.= rand(100) |> gpu

jvp(f,x,v) = ForwardDiff.partials.(f(ForwardDiff.Dual.(x, v)), 1)

function update!(opt, x, x̄::AbstractArray{<:ForwardDiff.Dual})
           x̄ = getindex.(ForwardDiff.partials.(x̄), 1)
           Flux.update!(opt, x, x̄)
end

loss(x,v) = sum(abs2,jvp(f,x,v))

loss(x,v) # Works

ps = Flux.params(f)
res,back = Zygote.pullback(ps) do
 loss(x,v)
end # Throws error here

update!(ADAM(0.01),ps,back(1f0))

```

---

<div class="post-metadata">

**Author:** ![mkalia94](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mkalia94/32/22849_2.png) [@mkalia94](https://discourse.julialang.org/u/mkalia94)\
**Post date:** [March 23, 2021, 12:38pm UTC](https://discourse.julialang.org/t/how-to-use-gradient-of-neural-network-as-the-loss-function/50569/11 "2021-03-23T12:38:33Z")

</div>

Perhaps @ChrisRackauckas could help out here?

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [March 23, 2021, 12:50pm UTC](https://discourse.julialang.org/t/how-to-use-gradient-of-neural-network-as-the-loss-function/50569/12 "2021-03-23T12:50:19Z")

</div>

It’s a Zygote issue. Use ReverseDiff over Zygote until Diffractor is released.

---

<div class="post-metadata">

**Author:** ![mkalia94](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mkalia94/32/22849_2.png) [@mkalia94](https://discourse.julialang.org/u/mkalia94)\
**Post date:** [March 23, 2021, 3:53pm UTC](https://discourse.julialang.org/t/how-to-use-gradient-of-neural-network-as-the-loss-function/50569/13 "2021-03-23T15:53:07Z")

</div>

So the code showed above will look a lot different, I observed that `ReverseDiff` does not have `partials`, and using `ReverseDiff.jacobian` instead leads to significantly more allocations (obviously) and a `this intrinsic must be compiled to be called` error…could you perhaps give a more explicit answer? Thanks!

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [March 23, 2021, 6:30pm UTC](https://discourse.julialang.org/t/how-to-use-gradient-of-neural-network-as-the-loss-function/50569/14 "2021-03-23T18:30:15Z")

</div>

You don’t need to modify much. There’s examples in the DiffEqFlux repo where the RHS of the ODE uses Zygote internally, and ReverseDiff for the VJP. It has to be that order though: the other fails.