Open arnauqb opened 1 year ago
Hey, apologies for the late reply. I've been taking a look into this, and it seems to be an interaction between Zygote's broadcast differentiation machinery and the fact that rand ∘ Bernoulli
produces a boolean. Things work fine for a geometric:
function sample_n_geometrics_vectorized(p, n)
probs = p * ones(n)
return sum(rand.(Geometric.(probs)))
end
Zygote.gradient(p -> sample_n_geometrics_vectorized(p, n), p) # works fine
So while this is not fixed, a couple alternatives you have are to rewrite using a map
or make sure the broadcasted function doesn't return a boolean:
function sample_n_bernoullis_vectorized_1(p, n)
probs = p * ones(n)
return (x -> rand(Bernoulli(x))*1).(probs) |> sum
end
function sample_n_bernoullis_vectorized_2(p, n)
probs = p * ones(n)
return map(rand ∘ Bernoulli, probs) |> sum
end
Zygote.gradient(p -> sample_n_bernoullis_vectorized_1(p, n), p) # works fine
Zygote.gradient(p -> sample_n_bernoullis_vectorized_2(p, n), p) # works fine
Thanks for catching this!
Hi there, first of all, thanks for this great package. I'm trying to understand how StochasticAD integrates with Zygote and I ran into a problem when trying to differentiate a "dot" operation. Here is the code:
Are vectorized operations not supported with Zygote?
Thanks!