What’s probably happening is that Zygote is trying to differentiate the code inside CUDA.randn
, which ultimately calls non-Julia code via llvmcall
.
The reason that it does not try to do this with randn
is that there’s a rule instructing it not to look, here.
You can define such a rule for CUDA.randn
in your code, or make a PR adding it for everyone here.
In general, you can also tell Zygote to ignore some bit of code by doing this (or the ChainRulesCore equivalent):
r = Zygote.@ignore CUDA.randn(T, 4,4)
r * v[1:4]