Hi! I’ve been running into memory/runtime issues with Zygote when using list comprehensions. I’d like to ask a simple question. Suppose
x is an array of length n. When Zygote differentiates backwards through the following code
f(x) = sum([x[i] - x[i+1] for i in range 1:2:N]),
does it use O(n^2) memory/time, or O(n)?
The argument for O(n^2) is that a list comprehension is essentially the application of the function
f(i) = x[i] - x[i+1] over
1:2:N. For each entry in the list comprehension, Zygote will compute the sensitivity of the output with respect to the variables that
f is a closure of (namely,
x), which has size O(n). This leads to a total of O(n^2) memory. Essentially, Zygote does not realize that the Jacobian of the list comprehension is sparse.
The argument for O(n) is that Zygote does actually realize this – for example, by representing \Delta f for each entry with O(1) memory using a sparse vector with two non zero entries.
I haven’t been able to parse through the Zygote code well enough to understand which of these cases is true. More generally, my question is: will Zygote always be smart enough to not use order of magnitudes more memory than necessary when differentiating through list comprehensions? Or should I be wary of this issue, and try to use maps, broadcasting, etc. whenever possible?