# Suggestions to improve Zygote performance for simple vector map/broadcast/comprehension?

**URL:** <https://discourse.julialang.org/t/suggestions-to-improve-zygote-performance-for-simple-vector-map-broadcast-comprehension/35207>\
**Category:** General Usage\
**Tags:** zygote\
**Created:** [February 27, 2020, 1:30am UTC](https://discourse.julialang.org/t/suggestions-to-improve-zygote-performance-for-simple-vector-map-broadcast-comprehension/35207 "2020-02-27T01:30:38Z")\
**Posts on this page:** 1\
**Showing post:** 5

<div class="post-metadata">

**Author:** ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)\
**Post date:** [November 9, 2022, 1:47am UTC](https://discourse.julialang.org/t/suggestions-to-improve-zygote-performance-for-simple-vector-map-broadcast-comprehension/35207/5 "2022-11-09T01:47:56Z")

</div>

It depends on what code is running inside the comprehension, but probably. The key performance sink pitfall in the OP is that it uses control flow (conditionals and loops). Zygote isn’t able to generate efficient code for functions using control flow, so you’ll see both slower speeds and more allocations. Array comprehensions/map/broadcast with these functions is a worst-case scenario because it literally multiplies the overhead over the number of elements processed.

We can show the impact of removing control flow by using a branchless conditional (`ifelse`) instead of the ternary:

```julia
build_vector2(x) = [ifelse(i<500, x, 0) for i=1:1000]

julia> @btime gradient(x -> sum(build_vector(x)), 1);
  830.159 μs (6559 allocations: 285.09 KiB)

julia> @btime gradient(x -> sum(build_vector2(x)), 1);
  17.263 μs (44 allocations: 119.25 KiB)

```

However, some functions must use control flow. In that case, you have a few options:

1. Use [API · ChainRules](https://juliadiff.org/ChainRulesCore.jl/stable/api.html#Ignoring-gradients) around functions/code blocks that use control flow but don’t need to be differentiated.
2. Define your own `rrule`(s) for functions that use control flow. The advice in [Writing good rules · ChainRules](https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/writing_good_rules.html) applies as always, but one additional concern here is to make sure the type of the returned _pullback function_ is stable. If it isn’t, you’ll run into many of the same issues as Zygote does.

---

_[View the full topic](https://discourse.julialang.org/t/suggestions-to-improve-zygote-performance-for-simple-vector-map-broadcast-comprehension/35207)._