Adding support for Where #117
Conversation
src/ops.jl
Outdated
| output[i] = y[i] | ||
| end | ||
| end | ||
| return output |
There was a problem hiding this comment.
Although this implementation works on CPU, most likely it won't (efficiently) work on GPU. How about using ifelse instead?
julia> ifelse.([true, false, true], [1, 1, 1], [2, 2, 2])
3-element Vector{Int64}:
1
2
1There was a problem hiding this comment.
I will go ahead and change this! Did my other implementations have similar issues? If so, should I go back and make sure there are no (or minimal) allocations in my implementations?
There was a problem hiding this comment.
No. Most Julia functions work equally well on CPU and GPU. Two things that are different for while are:
- Control flow on Julia level.
- Scalar indexing.
Both operations have bad performance on GPU due to their batch-first design.
Although not a problem for GPU per se, I'd also add to this list in-place update because they make quite some trouble for autodifferentiation frameworks.
|
Now it looks good. Thank you! |
Adding support for the Where layer of ONNX
PR Checklist