# How to make Zygote's gradient function work with a custom rrule

**URL:** <https://discourse.julialang.org/t/how-to-make-zygotes-gradient-function-work-with-a-custom-rrule/97439>\
**Category:** Machine Learning\
**Tags:** question\
**Created:** [April 13, 2023, 6:16pm UTC](https://discourse.julialang.org/t/how-to-make-zygotes-gradient-function-work-with-a-custom-rrule/97439 "2023-04-13T18:16:34Z")\
**Posts on this page:** 1\
**Page:** 1

<div class="post-metadata">

**Author:** ![gladisor](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gladisor/32/44475_2.png) [@gladisor](https://discourse.julialang.org/u/gladisor)\
**Post date:** [April 13, 2023, 6:16pm UTC](https://discourse.julialang.org/t/how-to-make-zygotes-gradient-function-work-with-a-custom-rrule/97439/1 "2023-04-13T18:16:35Z")

</div>

Here is a link to my issue which describes the problem I am facing:

> <https://github.com/JuliaDiff/ChainRulesCore.jl/issues/617#issue-1666784788>
>
> Hello all,
> 
> I am trying to extend the pedagogical example to work as if Foo is… a layer in the Flux ecosystem. I would like to insert Foo into a Flux Chain and train its parameters using the derivatives computed in the rrule instead of the standard Zygote AD. When I use the "gradient" function it returns a Grads(...) struct but the parameters do not match the gradients:
> 
> \`\`\`
> using Flux
> using Flux: params
> using ChainRulesCore
> 
> struct Foo
> A::Matrix
> c::Float64
> end
> 
> Flux.@functor Foo
> 
> function foo\_mul(foo::Foo, b::AbstractArray)
> return foo.A \* b
> end
> 
> function ChainRulesCore.rrule(::typeof(foo\_mul), foo::Foo, b::AbstractArray)
> println("calling foo rrule")
> y = foo\_mul(foo, b)
> 
> function foo\_mul\_pullback(ȳ)
> 
> f̄ = NoTangent()
> f̄oo = Tangent{Foo}(; A=ȳ \* b', c=ZeroTangent())
> b̄ = @thunk(foo.A' \* ȳ)
> 
> return f̄, f̄oo, b̄
> end
> 
> return y, foo\_mul\_pullback
> end
> 
> foo = Foo(randn(2, 2), 1.0)
> b = randn(2)
> 
> ps = params(foo)
> gs = gradient(() -\> sum(foo\_mul(foo, b)), ps)
> 
> display(gs.grads)
> 
> \`\`\`
> Here is the output of the code:
> 
> \`\`\`
> julia\> include("scripts/foo.jl")
> calling foo rrule
> IdDict{Any, Any} with 3 entries:
> :(Main.b) =\> \[1.5824, -0.280951\]
> \[0.795399 0.392306; 0.786999 -0.673257\] =\> nothing
> :(Main.foo) =\> (A = \[-0.728823 0.776745; -0.728823 0.776745\], c = nothing)
> \`\`\`
> 
> How can I make this work properly with the Flux "gradient" function?