Maybe I was unclear in that thread. The problem is control flow inside the body of the array comprehension and not with the comprehension itself. Notice how in Suggestions to improve Zygote performance for simple vector map/broadcast/comprehension? - #5 by ToucheSir I replace the ternary inside the [ ... for ... ]
with an ifelse
still inside the [ ... for ... ]
.
So you wouldn’t need a custom rule for compute_many_prods
, but one for (f,g) -> f * g * i
. Of course this particular example won’t be any faster because the f * g * i
doesn’t use any control flow, so you’ll want a better MWE that does.