# Manually giving gradients in Flux.jl

**URL:** <https://discourse.julialang.org/t/manually-giving-gradients-in-flux-jl/70459>\
**Category:** New to Julia\
**Tags:** flux, autodiff\
**Created:** [October 27, 2021, 7:58am UTC](https://discourse.julialang.org/t/manually-giving-gradients-in-flux-jl/70459 "2021-10-27T07:58:52Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![iHany](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ihany/32/18151_2.png) [@iHany](https://discourse.julialang.org/u/iHany)\
**Post date:** [October 27, 2021, 7:58am UTC](https://discourse.julialang.org/t/manually-giving-gradients-in-flux-jl/70459/1 "2021-10-27T07:58:52Z")

</div>

Hi,  
I’m solving a Q-learning problem.  
I’m trying to a gradient of TD error. That is,

TD-error = Q^{\theta}(x, u) - (r(x, u) + \min\_{v} Q^{\theta}(x^{+}, v))

where \theta = (\theta\_1, ..., \theta\_N) is the network parameter, x, u, and x^{+} are state, control, successor state.

The difficulty comes from \min. That is,

\frac{Q^{\theta}(x^{+}, u^{\*}(x^+, \theta))}{\partial \theta\_i} = \frac{\partial Q^{\theta}}{\partial \theta\_i} + \frac{\partial Q^{\theta}}{\partial u} \frac{\partial u^{\*}}{\partial \theta\_i}

I used to use Flux.jl to get a gradient automatically.  
The first term, \frac{\partial Q^{\theta}}{\partial \theta\_i} would be easy; just from

```julia
gs_theta = gradient(params(theta)) do
    Q(x_next, u_star, \theta)
end

```

as shown in [Flux.jl manual](https://fluxml.ai/Flux.jl/stable/models/basics/#Taking-Gradients).  
The second term would also be easy by using

```julia
gs_u = gradient(params(u_star)) do
    Q(x_next, u_star, \theta)
end

```

However, for the last term, the auto-calculation of the gradient of \min seems not to be supported by Flux.jl.  
It may be achievable by using [DiffOpt.jl](https://github.com/jump-dev/DiffOpt.jl), namely, `gs_u_star_theta_i` and the resulting gradient would be `gs_theta + gs_u * [gs_u_star_theta_1, ..., gs_u_star_theta_N]`.  
So I would like to manually tell that “the gradient is the above equation”.

My questions are:

1. How can I manually give a gradient in Flux.jl?
2. How can I merge gradients appropriately, for example, the `gs_u_star_theta_i`’s.

The above description would be poor. Please leave any comments 🙂

---

<div class="post-metadata">

**Author:** ![iiisoo](https://avatars.discourse-cdn.com/v4/letter/i/ecd19e/32.png) [@iiisoo](https://discourse.julialang.org/u/iiisoo)\
**Post date:** [October 27, 2021, 9:21am UTC](https://discourse.julialang.org/t/manually-giving-gradients-in-flux-jl/70459/2 "2021-10-27T09:21:00Z")

</div>

Your θ is the network parameter but where is your nn?

---

<div class="post-metadata">

**Author:** ![iHany](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ihany/32/18151_2.png) [@iHany](https://discourse.julialang.org/u/iHany)\
**Post date:** [October 27, 2021, 9:35am UTC](https://discourse.julialang.org/t/manually-giving-gradients-in-flux-jl/70459/3 "2021-10-27T09:35:34Z")

</div>

Q(x, u, theta) is itself a network.  
More precisely, Q is constructed as a function of nn(x, theta) and u in my case to make sure that Q is convex in u.

---

<div class="post-metadata">

**Author:** ![ianfiske](https://avatars.discourse-cdn.com/v4/letter/i/58f4c7/32.png) [@ianfiske](https://discourse.julialang.org/u/ianfiske)\
**Post date:** [October 27, 2021, 11:14am UTC](https://discourse.julialang.org/t/manually-giving-gradients-in-flux-jl/70459/4 "2021-10-27T11:14:58Z")

</div>

Flux relies on Zygote for autodiff. Zygote now uses ChainRules for the actual set of math rules to define the basic gradients transformations. The preferred way to define custom gradients now is through ChainRulesCore’s `rrule`:

[https://juliadiff.org/ChainRulesCore.jl/stable/#frule-and-rrule](https://juliadiff.org/ChainRulesCore.jl/stable/#frule-and-rrule)

(also see [Custom Adjoints · Zygote](https://fluxml.ai/Zygote.jl/latest/adjoints/) for the alternative legacy method and a bit more verbose explanation than what I just gave)