# Autograd in Flux

**URL:** <https://discourse.julialang.org/t/autograd-in-flux/56378>\
**Category:** General Usage\
**Created:** [March 3, 2021, 1:36am UTC](https://discourse.julialang.org/t/autograd-in-flux/56378 "2021-03-03T01:36:17Z")\
**Posts on this page:** 15\
**Page:** 1

<div class="post-metadata">

**Author:** ![Khemraj\_Shukla](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/khemraj_shukla/32/19917_2.png) [@Khemraj\_Shukla](https://discourse.julialang.org/u/Khemraj_Shukla)\
**Post date:** [March 3, 2021, 1:36am UTC](https://discourse.julialang.org/t/autograd-in-flux/56378/1 "2021-03-03T01:36:17Z")

</div>

Hello All-  
I need to compute the gradient of output of a neural network w.r.t input. In my case output is vector and input is also vector. I was trying to use `gradient` routine from `Flux` package but it was complaing that out put is not scaler.I will appreciate if can help me out with this.

Thank you very much!

Raj

---

<div class="post-metadata">

**Author:** ![asyrov](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/asyrov/32/13578_2.png) [@asyrov](https://discourse.julialang.org/u/asyrov)\
**Post date:** [March 3, 2021, 1:50am UTC](https://discourse.julialang.org/t/autograd-in-flux/56378/2 "2021-03-03T01:50:18Z")

</div>

I guess this may help:

> [@Computing function Jacobian from Zygote.pullback](https://discourse.julialang.org/t/computing-function-jacobian-from-zygote-pullback/42979):
>
> Zygote [documentation on custom adjoints](https://fluxml.ai/Zygote.jl/dev/adjoints/#Pullbacks-1) states dx/dy (in the notation from the docs) is the Jacobian when x, y are vectors. I am struggling in relating this to an implementation. For a function f : R^N -\> R^M modeled with f = Chain(Dense(N, H, tanh), Dense(H, M)) with hidden dimension H, I intuitively thought dy/dx would be an MxN matrix, but as we all know intuition often fails so I experimented a bit. Running using Flux, Zygote N, M, H = 2, 3, 16 g = Chain(Dense(N, H, tanh), Dense(H, M)) x …

---

<div class="post-metadata">

**Author:** ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)\
**Post date:** [March 3, 2021, 2:27am UTC](https://discourse.julialang.org/t/autograd-in-flux/56378/3 "2021-03-03T02:27:51Z")

</div>

Hello!

It’s difficult to diagnose the issue here without a MWE, can you read through the instructions in [Please read: make it easier to help you](https://discourse.julialang.org/t/psa-make-it-easier-to-help-you/14757) first?

---

<div class="post-metadata">

**Author:** ![Khemraj\_Shukla](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/khemraj_shukla/32/19917_2.png) [@Khemraj\_Shukla](https://discourse.julialang.org/u/Khemraj_Shukla)\
**Post date:** [March 3, 2021, 6:13am UTC](https://discourse.julialang.org/t/autograd-in-flux/56378/4 "2021-03-03T06:13:19Z")

</div>

Thanks! Sure  
I wrote following routine

```julia
function net_f(x, t, w, b; act=σ)
    u_nn = net_u(x, t, w, b; act=σ) # u_nn is of size (25600, 1)
    ## x, t, w and b are initialized as parameters
    Tracker.back!(u_nn, ones(size(x)[1]))
    u_x = Tracker.grad(x)
    u_t = Tracker.grad(t)
    return u_x, u_t
end

```

I need to compute the gradient of `u_nn` w.r.t. `t` and `x`, It is kind of computing partial derivative of output of neurl network wrt to input. In this implementation, I don’t know why if I am getting same values of `u_x` and `u_t`. I am not sure where I am doing the mistakes.I will appreciate your help.

Thanks!  
Raj

---

<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 3, 2021, 6:57am UTC](https://discourse.julialang.org/t/autograd-in-flux/56378/5 "2021-03-03T06:57:47Z")

</div>

This seems like you are using Flux with a Tracker.jl ad. Is this intentional? The current version of Flux uses Zygote AD, which has a different syntax.

Which version of Flux do you use?

---

<div class="post-metadata">

**Author:** ![Khemraj\_Shukla](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/khemraj_shukla/32/19917_2.png) [@Khemraj\_Shukla](https://discourse.julialang.org/u/Khemraj_Shukla)\
**Post date:** [March 3, 2021, 7:19am UTC](https://discourse.julialang.org/t/autograd-in-flux/56378/6 "2021-03-03T07:19:44Z")

</div>

@Tomas_Pevny Yes I am using Tracker.jl and using Flux.jl. I can update that to latest version. Thanks!

---

<div class="post-metadata">

**Author:** ![Khemraj\_Shukla](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/khemraj_shukla/32/19917_2.png) [@Khemraj\_Shukla](https://discourse.julialang.org/u/Khemraj_Shukla)\
**Post date:** [March 3, 2021, 7:21am UTC](https://discourse.julialang.org/t/autograd-in-flux/56378/7 "2021-03-03T07:21:12Z")

</div>

Also, The flux version 0.9.0.

---

<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 3, 2021, 7:30am UTC](https://discourse.julialang.org/t/autograd-in-flux/56378/8 "2021-03-03T07:30:09Z")

</div>

Any particular reason for using 0.9.0?

---

<div class="post-metadata">

**Author:** ![Khemraj\_Shukla](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/khemraj_shukla/32/19917_2.png) [@Khemraj\_Shukla](https://discourse.julialang.org/u/Khemraj_Shukla)\
**Post date:** [March 3, 2021, 7:36am UTC](https://discourse.julialang.org/t/autograd-in-flux/56378/9 "2021-03-03T07:36:42Z")

</div>

no not at all.I can update it.

---

<div class="post-metadata">

**Author:** ![Khemraj\_Shukla](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/khemraj_shukla/32/19917_2.png) [@Khemraj\_Shukla](https://discourse.julialang.org/u/Khemraj_Shukla)\
**Post date:** [March 5, 2021, 7:46am UTC](https://discourse.julialang.org/t/autograd-in-flux/56378/10 "2021-03-05T07:46:36Z")

</div>

@Tomas_Pevny Any pointers towards my question. Will really appreciate. Thanks!

---

<div class="post-metadata">

**Author:** ![roflmaostc](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/roflmaostc/32/30123_2.png) [@roflmaostc](https://discourse.julialang.org/u/roflmaostc)\
**Post date:** [March 5, 2021, 9:35am UTC](https://discourse.julialang.org/t/autograd-in-flux/56378/11 "2021-03-05T09:35:35Z")

</div>

As the error complains. The input for the gradient must be a scalar function (like a L2 loss eg.).

Did you have a look at the recent Flux documentation? I believe it’s well explained there.  
And as mentioned above, without a MWE it’s really hard to help you 🙄

---

<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 5, 2021, 10:11am UTC](https://discourse.julialang.org/t/autograd-in-flux/56378/12 "2021-03-05T10:11:34Z")

</div>

Please, add MWE.

---

<div class="post-metadata">

**Author:** ![Khemraj\_Shukla](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/khemraj_shukla/32/19917_2.png) [@Khemraj\_Shukla](https://discourse.julialang.org/u/Khemraj_Shukla)\
**Post date:** [March 5, 2021, 7:35pm UTC](https://discourse.julialang.org/t/autograd-in-flux/56378/13 "2021-03-05T19:35:06Z")

</div>

@roflmaostc @Tomas_Pevny Here is my entire code.base. Thank you for your help.

```julia-auto
using Revise
using MAT
using Flux
using Flux.Tracker
using Random
Random.seed!(1234)

push!(LOAD_PATH, "/Users/eklavya/WORK_RAJ/ML_JL/pinn_julia/src")
using pinn_subroutine
fh = matopen("/Users/eklavya/WORK_RAJ/ML_JL/pinn_julia/data/burgers_shock.mat")
x=read(fh, "x")
t=read(fh, "t")
X=x'.* ones(length(t))
T=ones(length(x))' .* t
xu = collect(Iterators.flatten(X))
tu = collect(Iterators.flatten(T))
uexact = real(read(fh, "usol"))
layers = [2, 20, 20, 1]
w,b = initialize_nn(layers)
X = hcat(xu, tu)
X = X'
u_nn = net_u(xu,tu, w, b; act=σ)

u_x, u_t = net_f(xu, tu, w, b; act=σ)

```

### Subroutines

```julia-auto
function net_f(x, t, w, b; act=σ)
    u_nn = net_u(x, t, w, b; act=σ) # u_nn is of size (25600, 1)
    ## x, t, w and b are initialized as parameters
    Tracker.back!(u_nn, ones(size(x)[1]))
    u_x = Tracker.grad(x)
    u_t = Tracker.grad(t)
    return u_x, u_t
end

function xavier_init(l)
    in_dim = l[1]
    out_dim = l[2]
    xavier_stddev = sqrt(2.0/(in_dim + out_dim))
    distMethod=rand(Truncated(Normal(0.0, xavier_stddev), -Inf, Inf))
    w=[distMethod for i in 1:in_dim, j in 1:out_dim]
    b=zeros(1, out_dim)
    return w, b
end

function initialize_nn(layer)
    w=Any[]
    b=Any[]
    for i in 1:length(layer)-1
        wl, bl= xavier_init([layer[i],layer[i+1]])
        push!(w, wl)
        push!(b, bl)
    end
    return w, b
end

function net_u(x, t, w, b; act=σ)
    X = hcat(x, t)
    for i in 1:length(w)-1
        Y = act.(X*w[i] .+ b[i])
        X=Y
    end
    Y = X*w[end] .+ b[end]
    return Y
end

```

---

<div class="post-metadata">

**Author:** ![roflmaostc](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/roflmaostc/32/30123_2.png) [@roflmaostc](https://discourse.julialang.org/u/roflmaostc)\
**Post date:** [March 5, 2021, 7:42pm UTC](https://discourse.julialang.org/t/autograd-in-flux/56378/14 "2021-03-05T19:42:30Z")

</div>

As mentioned above, you can only do a gradient of a scalar function. Second, at the moment we prefer Zygote as the AD for Flux. So you might prefer Zygote.

So what does `net_f` actually calculate?

It’s really hard to help here, because your problem is unclear (at least to me).

PS: By MWE we usually understand that one gets rid of unnecessary calculations and that one narrows down the problem to the core issues. I believe that many people don’t invest time to dig into several functions (which even include some data file loading) solving a complex problem.

---

<div class="post-metadata">

**Author:** ![Khemraj\_Shukla](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/khemraj_shukla/32/19917_2.png) [@Khemraj\_Shukla](https://discourse.julialang.org/u/Khemraj_Shukla)\
**Post date:** [March 5, 2021, 7:52pm UTC](https://discourse.julialang.org/t/autograd-in-flux/56378/15 "2021-03-05T19:52:47Z")

</div>

@roflmaostc Thank you! Yes actually `net_f` computes the gradient of output of network with respect to input. Input of network is multidimesional and output of network is a vector. Like the one offered by `tf.gradient` . I will prepare a brief working example and post it. Thank you for looking into it.

Thanks!