# Zygote calculate adjoint only for one of several arguments
**URL:** https://discourse.julialang.org/t/zygote-calculate-adjoint-only-for-one-of-several-arguments/59606
**Category:** Specific Domains
**Tags:** zygote, ad
**Created:** [April 19, 2021, 7:12pm UTC](https://discourse.julialang.org/t/zygote-calculate-adjoint-only-for-one-of-several-arguments/59606 "2021-04-19T19:12:59Z")
**Posts on this page:** 6
**Page:** 1
### Author:  [@roflmaostc](https://discourse.julialang.org/u/roflmaostc)
#### Post date: [April 19, 2021, 7:12pm UTC](https://discourse.julialang.org/t/zygote-calculate-adjoint-only-for-one-of-several-arguments/59606/1 "2021-04-19T19:12:59Z")
Hey,
is there a way that Zygote prevents the calculation of adjoints which are not needed at all?
Let’s say we have a function `f(a, b)` where we only want to have the first gradient:
```julia
b = #some value
g(a) = f(a, b)
Zygote.gradient(f, a)
```
Zygote now calculates (as far as I can see from my `rrule` definition) the adjoints with respect to `a` and `b` but the latter one is not needed at all and is actually wasted computing time.
Is there a way to prevent that?
> **Full Working Example**
>
> ```julia
> julia> function conv(u::AbstractArray{T, N}, v::AbstractArray{D, M}, dims=ntuple(+, max(N, M))) where {T, D, N, M}
> return real(ifft(fft(u, dims) .* fft(v, dims), dims))
> end
>
> julia> function ChainRulesCore.rrule(::typeof(conv), u::AbstractArray{T, N}, v::AbstractArray{D, M},
> dims=ntuple(+, max(N, M))) where {T, D, N, M}
> Y = conv(u, v, dims)
> function conv_pullback(barx)
> z = zero(eltype(u))
> return z, conv(barx, conj(v), dims), print("lo"), z
> end
> return Y, conv_pullback
> end
>
> julia> function main()
> u = randn((512, 512, 23))
> v = copy(u)
>
> f1(u) = sum(conv(u, v))
> f2(u, v) = sum(conv(u, v))
>
> @time f1(u)
> @time f1(u)
> @time f2(u, v)
> @time f2(u, v)
> @time Zygote.gradient(f1, u)
> @time Zygote.gradient(f1, u)
> @time Zygote.gradient(f1, u)
> @time Zygote.gradient(f2, u, v)
> @time Zygote.gradient(f2, u, v)
> @time Zygote.gradient(f2, u, v)
> return 0
> end
> main (generic function with 1 method)
>
> julia> main()
> 0.385862 seconds (106 allocations: 276.907 MiB, 6.92% gc time)
> 0.355024 seconds (106 allocations: 276.907 MiB)
> 0.499096 seconds (106 allocations: 276.907 MiB, 28.23% gc time)
> 0.368945 seconds (106 allocations: 276.907 MiB, 4.22% gc time)
> lo 0.830857 seconds (216 allocations: 599.815 MiB, 7.08% gc time)
> lo 0.871777 seconds (216 allocations: 599.815 MiB, 16.17% gc time)
> lo 0.760347 seconds (216 allocations: 599.815 MiB, 3.82% gc time)
> lo 0.804685 seconds (216 allocations: 599.815 MiB, 7.13% gc time)
> lo 0.912386 seconds (216 allocations: 599.815 MiB, 16.41% gc time)
> lo 0.804314 seconds (216 allocations: 599.815 MiB, 4.46% gc time)
>
> ```
Thanks,
Felix
---
### Author:  [@Tomas\_Pevny](https://discourse.julialang.org/u/Tomas_Pevny)
#### Post date: [April 20, 2021, 4:54am UTC](https://discourse.julialang.org/t/zygote-calculate-adjoint-only-for-one-of-several-arguments/59606/3 "2021-04-20T04:54:29Z")
I think it is not currently possible. There is a mcoro @thunk in diffrulescore, which should enable this, but it is disabled in zygote.
---
### Author:  [@Gattu\_Mytraya](https://discourse.julialang.org/u/Gattu_Mytraya)
#### Post date: [January 9, 2023, 2:32pm UTC](https://discourse.julialang.org/t/zygote-calculate-adjoint-only-for-one-of-several-arguments/59606/4 "2023-01-09T14:32:57Z")
Has there been any progress on this front?
---
### Author:  [@stevengj](https://discourse.julialang.org/u/stevengj)
#### Post date: [January 9, 2023, 3:02pm UTC](https://discourse.julialang.org/t/zygote-calculate-adjoint-only-for-one-of-several-arguments/59606/5 "2023-01-09T15:02:44Z")
> [@Tomas\_Pevny](#):
>
> I think it is not currently possible.
You can do it if you write a custom `rrule` for your function. Then you can return `@not_implemented("foo")` for the terms you don’t want.
Otherwise you are relying on the compiler to [eliminate “dead code”](https://en.wikipedia.org/wiki/Dead-code_elimination) for unused results.
In contrast, Enzyme.jl has separate explicit rules for constant arguments — see the post by @ChrisRackauckas contrasting Enzyme and Zygote rules: [What's the state of Automatic Differentiation in Julia January 2023? - #6 by ChrisRackauckas](https://discourse.julialang.org/t/whats-the-state-of-automatic-differentiation-in-julia-january-2023/92473/6)
---
### Author:  [@Gattu\_Mytraya](https://discourse.julialang.org/u/Gattu_Mytraya)
#### Post date: [January 9, 2023, 3:06pm UTC](https://discourse.julialang.org/t/zygote-calculate-adjoint-only-for-one-of-several-arguments/59606/6 "2023-01-09T15:06:24Z")
I’ll take a look, thanks!
---
### Author:  [@mcabbott](https://discourse.julialang.org/u/mcabbott)
#### Post date: [January 9, 2023, 3:06pm UTC](https://discourse.julialang.org/t/zygote-calculate-adjoint-only-for-one-of-several-arguments/59606/7 "2023-01-09T15:06:52Z")
There is [Zygote#966](https://github.com/FluxML/Zygote.jl/pull/966) which perhaps someone can push over the finish line.