Gradients of logpdf with TuringDiagMvNormal return nothing

[nmheim]

This puzzles me a bit

using DistributionsAD, Distributions, Flux
using DistributionsAD: TuringDiagMvNormal

Flux.@functor TuringDiagMvNormal

m = [1.0]
S = [0.1]

f = TuringDiagMvNormal(m,S)
x = 5 .+ randn(1,100)

ps = Flux.params(f)
#loss() = sum(logpdf(f,x))    # this returns nothings the loop below
loss() = loglikelihood(f,x)      # this works fine

gs = Flux.gradient(loss, ps)
for p in ps
    display(gs[p])
end

I think the logpdf is not hit during the gradient computation somehow. adding an error() in here still just returns nothings...

[mohamed82008]

This puzzles me a bit

Welcome to the club! It could be a Zygote bug.

[mohamed82008]

Try making a MWE

[nmheim]

I think this is related to FluxML/Zygote.jl#692 / FluxML/Zygote.jl#522, because this works:

loss(m,S) = sum(logpdf(TuringDiagMvNormal(m,S),x))
gs = Zygote.gradient(loss, m, S)