Re-using compiled code

I think that the problem is in pairwise function, which preallocates array for results and then fill it. It therefore performs a mutation, which is something Zygote cannot differentiate over.
You can either solved this by writing a rule, or by rewriting the pairwise such that it will not mutate existing code. For examle compute squared Euclidean distance would be

@. sum(x .^2 , dims = 1) - 2 * x' * y + sum(y .^2, dims = 1)' 
1 Like