gradient very slow when function call involves indexing

[benedict-96]

When I have an array in that is indexed in the function call in the first argument of gradient, the computation of the gradient is incredibly slow:

using Zygote, Printf, LinearAlgebra

number_data_points = 1000

data_input = Tuple([i] for i in 1:number_data_points)

function_to_be_differentiated(input, A) = norm(A*input)

function gradient_eval(num = Int(ceil(rand()*number_data_points)), A = rand(100000,1))
    input = data_input[num]
    @printf "First one: "
    @time Zygote.gradient(A -> function_to_be_differentiated(input, A), A)[1]
    @printf "Second one:"
    @time Zygote.gradient(A -> function_to_be_differentiated(data_input[num], A), A)[1]
    @printf "\n"
end

for i in 1:5
    gradient_eval()
end

This gives:

First one:   0.062275 seconds (320.72 k allocations: 19.683 MiB, 13.66% gc time)
Second one:  2.716012 seconds (10.85 M allocations: 694.051 MiB, 6.74% gc time)

First one:   0.000670 seconds (25 allocations: 2.290 MiB)
Second one:  0.963591 seconds (3.84 M allocations: 235.391 MiB, 9.53% gc time)

First one:   0.000627 seconds (25 allocations: 2.290 MiB)
Second one:  2.156478 seconds (6.43 M allocations: 413.277 MiB, 6.83% gc time)

First one:   0.000623 seconds (25 allocations: 2.290 MiB)
Second one:  2.993999 seconds (7.09 M allocations: 459.805 MiB, 6.93% gc time)

First one:   0.000676 seconds (25 allocations: 2.290 MiB)
Second one:  1.244626 seconds (3.16 M allocations: 190.261 MiB, 6.22% gc time)

Is this a known issue and if yes, can I expect to encounter similar performance limitations at different points?

[ToucheSir]

I would be curious as to why you're indexing a very large tuple with a dynamically defined, very large array. Even without Zygote, this should be quite slow because of how tuples work in Julia. If you must do this, then one thing to try would be to make data_input anything but a non-constant global: mark it const, define it with let, make it a parameter of gradient_eval or some outer function, etc.

[benedict-96]

Oh, I see. Thanks for the reply! If I do

const data_input = [[i] for i in 1:number_data_points

there is no such discrepancy in performance.