FFTs in probabilistic models

Specific Domains Probabilistic programming
,
2

sorry, the traceback is really long so I had to truncate some of it.

In contrast, explicitly constructing and using the DFT matrix works. I’d like to make it work with FFTW though.


N = 10
ℱ = fft(Matrix(1.0I, N, N), 1)

@model function naivefftmodel(x, y)
	mu ~ Normal(0, 1)
	q = ℱ * (mu .+ x)
	y ~ MvNormal(real.(q), 0.1)
	return y
end

x = rand(N)
y = fft(3.0 .+ x)
chain = sample(naivefftmodel(x, real.(y)) , NUTS(0.65), 1000)

Any pointers?

Zygote fails on this model also, but I think it’s unrelated. MWE:

Turing.setadbackend(:zygote)

@model function mymodel(y)
	x ~ filldist(Normal(0, 1), 10)
	y ~ MvNormal(x, 1)
end
y = rand(10)	
chain = sample(mymodel(y), NUTS(0.65), 1000)

the trace:

MethodError: no method matching (::ChainRulesCore.ProjectTo{Float64, NamedTuple{(), Tuple{}}})(::NamedTuple{(:x,), Tuple{Int64}})

Closest candidates are:

(::ChainRulesCore.ProjectTo{var"#s12", D} where {var"#s12"<:Real, D<:NamedTuple})(!Matched::Complex) at /Users/bld/.julia/packages/ChainRulesCore/MhCRV/src/projection.jl:186

(::ChainRulesCore.ProjectTo{var"#s12", D} where {var"#s12"<:Number, D<:NamedTuple})(!Matched::ChainRulesCore.Tangent{var"#s11", T} where {var"#s11"<:Number, T}) at /Users/bld/.julia/packages/ChainRulesCore/MhCRV/src/projection.jl:192

(::ChainRulesCore.ProjectTo{T, D} where D<:NamedTuple)(!Matched::ChainRulesCore.Tangent{var"#s12", T} where {var"#s12"<:T, T}) where T at /Users/bld/.julia/packages/ChainRulesCore/MhCRV/src/projection.jl:142

…

  1. (::ChainRulesCore.ProjectTo{Ref, NamedTuple{(:type, :x), Tuple{DataType, ChainRulesCore.ProjectTo{Float64, NamedTuple{(), Tuple{}}}}}})(::Base.RefValue{Any})@ projection.jl:284
  2. #56 @ projection.jl:237 [inlined]
  3. #4 @ generator.jl:36 [inlined]
  4. iterate @ generator.jl:47 [inlined]
  5. collect (::Base.Generator{Base.Iterators.Zip{Tuple{Vector{ChainRulesCore.ProjectTo{Ref, NamedTuple{(:type, :x), Tuple{DataType, ChainRulesCore.ProjectTo{Float64, NamedTuple{(), Tuple{}}}}}}}, Vector{Base.RefValue{Any}}}}, Base.var"#4#5"{ChainRulesCore.var"#56#57"}})@ array.jl:678
  6. map @ abstractarray.jl:2383 [inlined]
  7. (::ChainRulesCore.ProjectTo{AbstractArray, NamedTuple{(:elements, :axes), Tuple{Vector{ChainRulesCore.ProjectTo{Ref, NamedTuple{(:type, :x), Tuple{DataType, ChainRulesCore.ProjectTo{Float64, NamedTuple{(), Tuple{}}}}}}}, Tuple{Base.OneTo{Int64}}}}})(::Vector{Base.RefValue{Any}})@ projection.jl:237
  8. (::ChainRules.var"#sum_pullback#1375"{Colon, typeof(getindex), ChainRulesCore.ProjectTo{AbstractArray, NamedTuple{(:elements, :axes), Tuple{Vector{ChainRulesCore.ProjectTo{Ref, NamedTuple{(:type, :x), Tuple{DataType, ChainRulesCore.ProjectTo{Float64, NamedTuple{(), Tuple{}}}}}}}, Tuple{Base.OneTo{Int64}}}}}, Vector{Zygote.var"#ad_pullback#45"{Tuple{typeof(getindex), Base.RefValue{Float64}}, typeof(∂(getindex))}}})(::Int64)@ mapreduce.jl:88
  9. ZBack @ chainrules.jl:168 [inlined]
  10. Pullback @ threadsafe.jl:25 [inlined]
  11. (::typeof(∂(getlogp)))(::Int64)@ interface2.jl:0
  12. Pullback @ model.jl:439 [inlined]
  13. (::typeof(∂(evaluate_threadsafe)))(::Nothing)@ interface2.jl:0
  14. Pullback @ model.jl:391 [inlined]
  15. (::typeof(∂(λ)))(::Nothing)@ interface2.jl:0
  16. Pullback @ model.jl:383 [inlined]
  17. (::typeof(∂(λ)))(::Nothing)@ interface2.jl:0
  18. #203 @ lib.jl:203 [inlined]
  19. #1734#back @ adjoint.jl:67 [inlined]
  20. Pullback @ model.jl:396 [inlined]
  21. (::typeof(∂(λ)))(::Nothing)@ interface2.jl:0
  22. Pullback @ ad.jl:165 [inlined]
  23. (::typeof(∂(λ)))(::Int64)@ interface2.jl:0
  24. (::Zygote.var"#50#51"{typeof(∂(λ))})(::Int64)@ interface.jl:41
  25. gradient_logp (::Turing.Core.ZygoteAD, ::Vector{Float64}, ::DynamicPPL.TypedVarInfo{NamedTuple{(:x,), Tuple{DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:x, Tuple{}}, Int64}, Vector{DistributionsAD.TuringScalMvNormal{Vector{Float64}, Float64}}, Vector{AbstractPPL.VarName{:x, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}}, Float64}, ::DynamicPPL.Model{Main.workspace388.var"#mymodel#1", (:y,), (), (), Tuple{Vector{Float64}}, Tuple{}, DynamicPPL.DefaultContext}, ::DynamicPPL.Sampler{Turing.Inference.NUTS{Turing.Core.ZygoteAD, (), AdvancedHMC.DiagEuclideanMetric}}, ::DynamicPPL.DefaultContext)@ ad.jl:171
  26. gradient_logp @ ad.jl:83 [inlined]
  27. ∂logπ∂θ @ hmc.jl:433 [inlined]
  28. ∂H∂θ @ hamiltonian.jl:31 [inlined]
  29. phasepoint @ hamiltonian.jl:76 [inlined]
  30. phasepoint (::Random._GLOBAL_RNG, ::Vector{Float64}, ::AdvancedHMC.Hamiltonian{AdvancedHMC.DiagEuclideanMetric{Float64, Vector{Float64}}, Turing.Inference.var"#logπ#52"{DynamicPPL.TypedVarInfo{NamedTuple{(:x,), Tuple{DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:x, Tuple{}}, Int64}, Vector{DistributionsAD.TuringScalMvNormal{Vector{Float64}, Float64}}, Vector{AbstractPPL.VarName{:x, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}}, Float64}, DynamicPPL.Sampler{Turing.Inference.NUTS{Turing.Core.ZygoteAD, (), AdvancedHMC.DiagEuclideanMetric}}, DynamicPPL.Model{Main.workspace388.var"#mymodel#1", (:y,), (), (), Tuple{Vector{Float64}}, Tuple{}, DynamicPPL.DefaultContext}}, Turing.Inference.var"#∂logπ∂θ#51"{DynamicPPL.TypedVarInfo{NamedTuple{(:x,), Tuple{DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:x, Tuple{}}, Int64}, Vector{DistributionsAD.TuringScalMvNormal{Vector{Float64}, Float64}}, Vector{AbstractPPL.VarName{:x, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}}, Float64}, DynamicPPL.Sampler{Turing.Inference.NUTS{Turing.Core.ZygoteAD, (), AdvancedHMC.DiagEuclideanMetric}}, DynamicPPL.Model{Main.workspace388.var"#mymodel#1", (:y,), (), (), Tuple{Vector{Float64}}, Tuple{}, DynamicPPL.DefaultContext}}})@ hamiltonian.jl:153
  31. var"#initialstep#41" (::Nothing, ::Int64, ::Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, ::typeof(DynamicPPL.initialstep), ::Random._GLOBAL_RNG, ::DynamicPPL.Model{Main.workspace388.var"#mymodel#1", (:y,), (), (), Tuple{Vector{Float64}}, Tuple{}, DynamicPPL.DefaultContext}, ::DynamicPPL.Sampler{Turing.Inference.NUTS{Turing.Core.ZygoteAD, (), AdvancedHMC.DiagEuclideanMetric}}, ::DynamicPPL.TypedVarInfo{NamedTuple{(:x,), Tuple{DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:x, Tuple{}}, Int64}, Vector{DistributionsAD.TuringScalMvNormal{Vector{Float64}, Float64}}, Vector{AbstractPPL.VarName{:x, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}}, Float64})@ hmc.jl:167
  32. var"#step#17" (::Nothing, ::Base.Iterators.Pairs{Symbol, Int64, Tuple{Symbol}, NamedTuple{(:nadapts,), Tuple{Int64}}}, ::typeof(AbstractMCMC.step), ::Random._GLOBAL_RNG, ::DynamicPPL.Model{Main.workspace388.var"#mymodel#1", (:y,), (), (), Tuple{Vector{Float64}}, Tuple{}, DynamicPPL.DefaultContext}, ::DynamicPPL.Sampler{Turing.Inference.NUTS{Turing.Core.ZygoteAD, (), AdvancedHMC.DiagEuclideanMetric}})@ sampler.jl:87
  33. macro expansion @ sample.jl:123 [inlined]
  34. macro expansion @ ProgressLogging.jl:328 [inlined]
  35. macro expansion @ logging.jl:8 [inlined]
  36. var"#mcmcsample#20" (::Bool, ::String, ::Nothing, ::Int64, ::Int64, ::Type, ::Base.Iterators.Pairs{Symbol, Int64, Tuple{Symbol}, NamedTuple{(:nadapts,), Tuple{Int64}}}, ::typeof(AbstractMCMC.mcmcsample), ::Random._GLOBAL_RNG, ::DynamicPPL.Model{Main.workspace388.var"#mymodel#1", (:y,), (), (), Tuple{Vector{Float64}}, Tuple{}, DynamicPPL.DefaultContext}, ::DynamicPPL.Sampler{Turing.Inference.NUTS{Turing.Core.ZygoteAD, (), AdvancedHMC.DiagEuclideanMetric}}, ::Int64)@ sample.jl:114
  37. var"#sample#40" (::Type, ::Nothing, ::Bool, ::Int64, ::Bool, ::Int64, ::Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, ::typeof(StatsBase.sample), ::Random._GLOBAL_RNG, ::DynamicPPL.Model{Main.workspace388.var"#mymodel#1", (:y,), (), (), Tuple{Vector{Float64}}, Tuple{}, DynamicPPL.DefaultContext}, ::DynamicPPL.Sampler{Turing.Inference.NUTS{Turing.Core.ZygoteAD, (), AdvancedHMC.DiagEuclideanMetric}}, ::Int64)@ hmc.jl:133
  38. sample @ hmc.jl:116 [inlined]
  39. #sample#2 @ Inference.jl:142 [inlined]
  40. sample @ Inference.jl:142 [inlined]
  41. #sample#1 @ Inference.jl:132 [inlined]
  42. sample (::DynamicPPL.Model{Main.workspace388.var"#mymodel#1", (:y,), (), (), Tuple{Vector{Float64}}, Tuple{}, DynamicPPL.DefaultContext}, ::Turing.Inference.NUTS{Turing.Core.ZygoteAD, (), AdvancedHMC.DiagEuclideanMetric}, ::Int64)@ Inference.jl:132
  43. top-level scope @ *[Local: 11]

version info:

Julia Version 1.6.0
Commit f9720dc2eb (2021-03-24 12:55 UTC)
Platform Info:
OS: macOS (x86_64-apple-darwin19.6.0)
CPU: Intel(R) Core™ i7-4770HQ CPU @ 2.20GHz
WORD_SIZE: 64
LIBM: libopenlibm
LLVM: libLLVM-11.0.1 (ORCJIT, haswell)
Environment:
JULIA_NUM_THREADS = 4
JULIA_REVISE_WORKER_ONLY = 1