Help using CUDA, Zygote, and random numbers

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]
1 Like