# Zygote gradient results in slow Tuple getindex calls

**URL:** <https://discourse.julialang.org/t/zygote-gradient-results-in-slow-tuple-getindex-calls/48767>\
**Category:** Performance\
**Tags:** zygote\
**Created:** [October 21, 2020, 6:49pm UTC](https://discourse.julialang.org/t/zygote-gradient-results-in-slow-tuple-getindex-calls/48767 "2020-10-21T18:49:33Z")\
**Posts on this page:** 7\
**Page:** 1

<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:** [October 21, 2020, 6:49pm UTC](https://discourse.julialang.org/t/zygote-gradient-results-in-slow-tuple-getindex-calls/48767/1 "2020-10-21T18:49:34Z")

</div>

Hey,

I’m working on an optimization problem and I’m using Zygote to calculate the gradients ([code](https://github.com/roflmaostc/DeconvOptim.jl)).

Initially a lot of the chain rules were implemented manually by code but then I switched so that Zygote can do this stuff.  
Unfortunately the performance seemed to drop despite I defined computationally intensive adjoints explicitly. So I would have expected that Zygote only has to apply a few chain rules and summations.

To find the source of the performance drop, I did a profiling.

 ![Screenshot_2020-10-21_20-37-38](https://global.discourse-cdn.com/julialang/original/3X/c/9/c98dbba5559eda4eda9daf8f6e35f6edbf961945.png)

That `fft.jl` uses a lot of time is reasonable since all computationally expensive calculations are FFTs. I’m more wondering about `tuple.jl` and `getindex`. By clicking on `getindex` I got information that it is also connected to tuples.

In my source code I’m not using tuples directly except for the adjoint definitions. For example

```julia
function conv_aux(conv, rec, otf)
    return conv(rec, otf)
end

function ChainRulesCore.rrule(::typeof(conv_aux), conv, rec, otf)
    Y = conv_aux(conv, rec, otf)
    function conv_aux_pullback(barx)
        return NO_FIELDS, DoesNotExist(), conv(barx, conj(otf)), DoesNotExist()
    end
    return Y, conv_aux_pullback
end

```

and [here](https://github.com/roflmaostc/DeconvOptim.jl/blob/master/src/lossfunctions.jl).

I’m not really experienced in profiling or using AD, and I would be really happy if someone could give me some hints how to improve performance and why the tuples seem to cause issues.

Thanks a lot,

Felix

---

<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:** [October 21, 2020, 8:47pm UTC](https://discourse.julialang.org/t/zygote-gradient-results-in-slow-tuple-getindex-calls/48767/2 "2020-10-21T20:47:00Z")

</div>

I guess I found the reason for the many calls of tuple.  
It was connected to a `max(0.1, conv_aux(conv, rec, otf)`. I could reproduce that in a smaller example.

From a mathematical point of view that’s reasonable to cause troubles because this `max` function is not differentiable at 0.1.  
I still would be interested why Zygote ends up in that tuples calls a lot.

Thanks and sorry that I didn’t really provide a small standalone example,

Felix

---

<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:** [October 22, 2020, 8:22am UTC](https://discourse.julialang.org/t/zygote-gradient-results-in-slow-tuple-getindex-calls/48767/3 "2020-10-22T08:22:08Z")

</div>

Hey,

I can now provide a minimal working example:

```julia
using Zygote
using FFTW
using Profile
using StatProfilerHTML

N =1000
global psf = randn((N, N))
global img = randn((N, N))

function conv(img, psf)
    f_psf = fft(psf)
    f_img = fft(img)
    conv_res = real(ifft(f_psf .* f_img))
    #return max.(0.0, conv_res)
    return conv_res    
end 

f(img) = sum(conv(img, psf));

Zygote.gradient(f, img)[1]
@profilehtml Zygote.gradient(f, img)[1]

```

The results with `max.`:

 ![with_max](https://global.discourse-cdn.com/julialang/original/3X/4/0/40a7d870d0a52846214e979d05e6a012902b7c4b.png)

And without:

 ![Screenshot_2020-10-22_10-19-45](https://global.discourse-cdn.com/julialang/original/3X/6/5/658fe8ee967cb2cf334767d7752d6aabb3a6c9fa.png)

So I’m not sure whether I’m at the right place to discuss this.

Thanks,

Felix

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [October 22, 2020, 8:50am UTC](https://discourse.julialang.org/t/zygote-gradient-results-in-slow-tuple-getindex-calls/48767/4 "2020-10-22T08:50:16Z")

</div>

I think what you’re seeing here is that Zygote’s broadcasting isn’t very fast, except on cases where it has a special rule for that function:

```julia
julia> @btime gradient(x -> sum(1 .+ max.(0.0,x)), mat) setup=(mat=randn(100,100));
  990.097 μs (50054 allocations: 1.84 MiB)

julia> using Flux # Flux.relu has a special rule

julia> @btime gradient(x -> sum(1 .+ relu.(x)), mat) setup=(mat=randn(100,100));
  17.888 μs (6 allocations: 234.61 KiB)

julia> using Tullio # handles gradients without Broadcast.

julia> @btime gradient(x -> sum(@tullio y[i,j] := 1 + max(0.0,x[i,j])), mat) setup=(mat=randn(100,100));
  22.290 μs (18 allocations: 156.80 KiB)

julia> using Tracker # different broadcasting, perhaps using ForwardDiff

julia> @btime Tracker.gradient(x -> sum(1 .+ max.(0.0,x)), mat) setup=(mat=randn(100,100));
  94.865 μs (173 allocations: 709.87 KiB)

```

---

<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:** [October 22, 2020, 8:59am UTC](https://discourse.julialang.org/t/zygote-gradient-results-in-slow-tuple-getindex-calls/48767/5 "2020-10-22T08:59:49Z")

</div>

Hey,

thanks for the details.

Is it `max.(0.0, x)` or `1 .+ ...` which causes the problems? I guess the first one, right?

Thanks,

Felix

---

<div class="post-metadata">

**Author:** ![jonasaugust](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/jonasaugust/32/11942_2.png) [@jonasaugust](https://discourse.julialang.org/u/jonasaugust)\
**Post date:** [May 20, 2021, 6:16pm UTC](https://discourse.julialang.org/t/zygote-gradient-results-in-slow-tuple-getindex-calls/48767/6 "2021-05-20T18:16:25Z")

</div>

Notice the crazy number of allocations in the first of your examples: about 5 per element of a 10000 entry array. The other non-Zygote examples have much less allocations.

---

<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:** [May 20, 2021, 6:19pm UTC](https://discourse.julialang.org/t/zygote-gradient-results-in-slow-tuple-getindex-calls/48767/7 "2021-05-20T18:19:13Z")

</div>

I often observe a lot of allocations with Zygote and so far I haven’t really figured how to fix that properly.

Usually the forward passes are type stable in these situations.