# Rrule for zygote

**URL:** <https://discourse.julialang.org/t/rrule-for-zygote/131069>\
**Category:** New to Julia\
**Tags:** zygote, gridap, chainrulescore, gradient, chainrules\
**Created:** [July 25, 2025, 10:07pm UTC](https://discourse.julialang.org/t/rrule-for-zygote/131069 "2025-07-25T22:07:19Z")\
**Posts on this page:** 8\
**Page:** 1

<div class="post-metadata">

**Author:** ![mary](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mary/32/53139_2.png) [@mary](https://discourse.julialang.org/u/mary)\
**Post date:** [July 25, 2025, 10:07pm UTC](https://discourse.julialang.org/t/rrule-for-zygote/131069/1 "2025-07-25T22:07:19Z")

</div>

Hi, I am using rrule for gradient calculation in zygote. I defined a function which calculates the gradients and I just call that function inside the entire gradient calculation. Can zygote handle this type of rule?

```julia-auto
function ldt_simulation(nsamples,iteration_counter,pf_vec)
.
.
.
    return grad_logPf 
end

```

```julia-auto
function Dgfdpf(pf_vec; iteration_counter,β, η, fem_params)
    pfg = ldt_simulation(nsamples,iteration_counter,pf_vec)
    return pfg
end 

function pf_p0(p0; r, fem_params)
    pf_vec = Filter(p0; r, fem_params)
    pf_vec
end

function rrule(::typeof(pf_p0), p0; r, fem_params)
  function pf_pullback(dgdpf)
    NO_FIELDS, Dgdp(dgdpf; r, fem_params)
  end
  pf_p0(p0; r, fem_params), pf_pullback
end

function Dgdp(dgdpf; r, fem_params)
    Af = assemble_matrix(fem_params.Pf, fem_params.Qf) do u, v
        ∫(a_f(r, u, v))fem_params.dΩ + ∫(v * u)fem_params.dΩ
    end
    wvec = Af' \ dgdpf
    wh = FEFunction(fem_params.Pf, wvec)
    l_temp(dp) = ∫(wh * dp)fem_params.dΩ
    return assemble_vector(l_temp, fem_params.P)
end 

function gf_p(p0::Vector;iteration_counter, r, β, η, fem_params)
    pf_vec = pf_p0(p0; r, fem_params)
    gf_pf(pf_vec;iteration_counter, β, η, fem_params)
end

function gf_p(p0::Vector, grad::Vector;iteration_counter, r, β, η, fem_params)
    if length(grad) > 0
        
        dgdp, = Zygote.gradient(p -> gf_p(p; iteration_counter,r, β, η, fem_params), p0)
        grad[:] = dgdp
    end
    gvalue = gf_p(p0::Vector; iteration_counter,r, β, η, fem_params)
    
    gvalue
end

```

Thank you

---

<div class="post-metadata">

**Author:** ![stevengj](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/stevengj/32/71_2.png) [@stevengj](https://discourse.julialang.org/u/stevengj)\
**Post date:** [July 25, 2025, 11:26pm UTC](https://discourse.julialang.org/t/rrule-for-zygote/131069/2 "2025-07-25T23:26:35Z")

</div>

> [@mary](#):
>
> I defined a function which calculates the gradients and I just call that function inside the entire gradient calculation. Can zygote handle this type of rule?

Not quite. An `rrule` is not just a function that computes a gradient, it is a vector–Jacobian product.

That is, if you have a function f(x), then an `rrule` for `x` returns a “pullback” function `g(v)` that computes v^T f'(x) where f' is the Jacobian.

If you have a scalar-valued function f(x) of a vector x, then the Jacobian is a row vector f'(x) = (\nabla f)^T, and your pullback takes a _scalar_ v and returns v (\nabla f)^T (in the correct shape).

---

<div class="post-metadata">

**Author:** ![mary](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mary/32/53139_2.png) [@mary](https://discourse.julialang.org/u/mary)\
**Post date:** [July 26, 2025, 3:27am UTC](https://discourse.julialang.org/t/rrule-for-zygote/131069/3 "2025-07-26T03:27:59Z")

</div>

I have a function`f(x) = K[(ln(pf)-ln(pa))^+]^2` which has a scalar value if inside the bracket has a + value or zero if inside the bracket isn’t positive. The `ldt_simulation` function calculates the gradient of this `f(x)` which is a vector. Based on this assumptions and what you explained I think it can handle the gradient calculation by zygote, Am I right? Is there any way to check this?

---

<div class="post-metadata">

**Author:** ![stevengj](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/stevengj/32/71_2.png) [@stevengj](https://discourse.julialang.org/u/stevengj)\
**Post date:** [July 26, 2025, 12:21pm UTC](https://discourse.julialang.org/t/rrule-for-zygote/131069/4 "2025-07-26T12:21:53Z")

</div>

> [@mary](#):
>
> Based on this assumptions and what you explained I think it can handle the gradient calculation by zygote, Am I right? Is there any way to check this?

Yes, you can implement an `rrule` for anything (assuming the derivative exists). You normally check it with finite difference.

---

<div class="post-metadata">

**Author:** ![marcobonici](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/marcobonici/32/20549_2.png) [@marcobonici](https://discourse.julialang.org/u/marcobonici)\
**Post date:** [July 26, 2025, 2:10pm UTC](https://discourse.julialang.org/t/rrule-for-zygote/131069/5 "2025-07-26T14:10:44Z")

</div>

Often I am able to check with ForwardDiff.jl as well (maybe less performant for high number of inputs, but for sure easy to use in a lot of case).

---

<div class="post-metadata">

**Author:** ![stevengj](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/stevengj/32/71_2.png) [@stevengj](https://discourse.julialang.org/u/stevengj)\
**Post date:** [July 26, 2025, 4:14pm UTC](https://discourse.julialang.org/t/rrule-for-zygote/131069/6 "2025-07-26T16:14:48Z")

</div>

I think the OP @mary is trying to do [topology optimization](https://en.wikipedia.org/wiki/Topology_optimization) through Gridap.jl FEM, similar to [this tutorial](https://gridap.github.io/Tutorials/dev/pages/t019_TopOptEMFocus/), in which case ForwardDiff may not be so easy.

I would usually finite difference to compare the gradient to some random directional derivatives, i.e. check that f(x + \delta x) - f(x) \approx \delta x^T \nabla f for some small random vectors \delta x.

---

<div class="post-metadata">

**Author:** ![mary](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mary/32/53139_2.png) [@mary](https://discourse.julialang.org/u/mary)\
**Post date:** [July 29, 2025, 12:56am UTC](https://discourse.julialang.org/t/rrule-for-zygote/131069/7 "2025-07-29T00:56:14Z")

</div>

I wanted to check that with finite difference but the objective is K[(ln(pf)-ln(pa))^+]^2 which pf is probability of failure. By perturbing a small value the pf is not changing so the finite difference value is 0.0 but I have a huge gradient because in addition to 2K[(ln(pf)-ln(pa))^+] I have the \nabla ln(pf) multiplied. I am not sure that finite difference is working here or not.

---

<div class="post-metadata">

**Author:** ![stevengj](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/stevengj/32/71_2.png) [@stevengj](https://discourse.julialang.org/u/stevengj)\
**Post date:** [July 29, 2025, 1:37am UTC](https://discourse.julialang.org/t/rrule-for-zygote/131069/8 "2025-07-29T01:37:23Z")

</div>

> [@mary](#):
>
> By perturbing a small value the pf is not changing so the finite difference value is 0.0

If you perturb the inputs by too little then the accuracy of finite differences is destroyed by the finite precision. See also [chapter 4 of our matrix-calculus notes](https://arxiv.org/abs/2501.14787).