DynamicPPL.jl icon indicating copy to clipboard operation
DynamicPPL.jl copied to clipboard

`eltype` for empty `VarInfo` always results in `Any`

Open torfjelde opened this issue 3 years ago • 12 comments

The bug and the fix:

julia> using Revise, DynamicPPL, Distributions

julia> @model demo() = x ~ Normal()
demo (generic function with 2 methods)

julia> model = demo() | (x = 1.0, )
Model{typeof(demo), (), (), (), Tuple{}, Tuple{}, ConditionContext{(:x,), NamedTuple{(:x,), Tuple{Float64}}, DefaultContext}}(:demo, demo, NamedTuple(), NamedTuple(), ConditionContext((x = 1.0,), DefaultContext()))

julia> varinfo = VarInfo(model)
TypedVarInfo{NamedTuple{(), Tuple{}}, Float64}(NamedTuple(), Base.RefValue{Float64}(-1.4189385332046727), Base.RefValue{Int64}(1))

julia> eltype(varinfo, SampleFromPrior())
Any

julia> # Overload "empty" `TypedVarInfo`.
       function Base.eltype(
           vi::DynamicPPL.VarInfo{NamedTuple{(),Tuple{}}},
           spl::Union{DynamicPPL.AbstractSampler,DynamicPPL.SampleFromPrior}
       )
           return eltype(DynamicPPL.getlogp(vi))
       end

julia> eltype(varinfo, SampleFromPrior())
Float64

This can be pretty annoying as it leads to downstream issues such as StackOverflowError when used with MvNormal.

torfjelde avatar Mar 01 '22 23:03 torfjelde

Ǹot really surprising that sometimes Any is returned - seems we use Core.Compiler.return_type: https://github.com/TuringLang/DynamicPPL.jl/blob/083dfa11447e762aaa74a07e08c39f40b39a7d0e/src/varinfo.jl#L1036 It is allowed to return Any (or any other super type of the actual eltype) for any arguments, not only empty VarInfos :shrug:

devmotion avatar Mar 01 '22 23:03 devmotion

Yep:) This is a thing we keep coming back to :sweat_smile:

torfjelde avatar Mar 01 '22 23:03 torfjelde

Yeah :smile: My main point is that there might be many more cases where Any (or some other super type) is returned, so there's no guarantee that the fix will actually fix the problem generally. And I don't think this can be avoided as long as we use Core.Compiler.return_type.

devmotion avatar Mar 01 '22 23:03 devmotion

BTW the fix is inconsistent with the behaviour for non-empty VarInfo:

julia> using DynamicPPL, Distributions, Random

julia> @model demo() = x ~ Normal()
demo (generic function with 2 methods)

julia> eltype(VarInfo(demo()), SampleFromPrior())
Float64

julia> typeof(getlogp(VarInfo(demo())))
Float64

julia> function varinfo_float32(model)
           vi = VarInfo(DynamicPPL.Metadata(), Ref{Float32}(0.0), Ref(0))
           model(vi)
           return TypedVarInfo(vi)
       end
varinfo_float32 (generic function with 1 method)

julia> eltype(varinfo_float32(demo()), SampleFromPrior())
Float64

julia> typeof(getlogp(varinfo_float32(demo())))
Float32

I.e., usually the eltype and the type of getlogp are not the same.

devmotion avatar Mar 01 '22 23:03 devmotion

And I don't think this can be avoided as long as we use Core.Compiler.return_type.

Agreed :+1:

BTW the fix is inconsistent with the behaviour for non-empty VarInfo

Sure! But in this case I'd say that this fix is still better than not :sweat_smile:

torfjelde avatar Mar 01 '22 23:03 torfjelde

Sure! But in this case I'd say that this fix is still better than not sweat_smile

But why base it on getlogp if it's not based on getlogp if the VarInfo is not empty? Why not just always return Float64?

devmotion avatar Mar 02 '22 00:03 devmotion

But why base it on getlogp if it's not based on getlogp if the VarInfo is not empty? Why not just always return Float64?

Sure :+1:

Though IMO it should be based on getlogp. This is what we do for SimpleVarInfo. It means that we no longer need to make use of return_type, etc.

torfjelde avatar Mar 02 '22 00:03 torfjelde

This is what we do for SimpleVarInfo

That's definitely inconsistent with VarInfo :disappointed: Do we have to define eltype for VarInfo at all? It's not completely clear on what it should be based - but intuitively I would have assumed it should depend on the type of the samples since it is a datastructure for samples. I guess due ot the linearization/vectorization in VarInfo it was a bit easier to talk about the eltype of the samples.

devmotion avatar Mar 02 '22 00:03 devmotion

That's definitely inconsistent with VarInfo disappointed Do we have to define eltype for VarInfo at all?

Yup! But there are other things which are also inconsistent with VarInfo..

As things are right now we unfortunately need to define eltype because we use this extensively to replace these ::Type parameters.

torfjelde avatar Mar 02 '22 00:03 torfjelde

Maybe we can fix these eltype use cases. But for this use case it seems more reasonable to base eltype on the samples but not the logp since the ::Type parameters are used e.g. to allocate arrays of samples - but I haven't seen an example where they were used to instantiate a log density calculation.

devmotion avatar Mar 02 '22 00:03 devmotion

A brief comment on the inconsistency between VarInfo / SimpleVarInfo: we would like to depreciate VarInfo gradually, so these inconsistencies won't matter that much in the long term.

yebai avatar Mar 02 '22 09:03 yebai

Yes, I guess probably this inconsistency between VarInfo and SimpleVarInfo is the least important issue.

I'm more concerned with defining eltype for varinfos at all, and particularly based on the log density. I think it's not completely obvious what eltype should be for such a datastructure of samples, in particular if the individual samples have different types. But intuitively I would not base it on the log density, and such a choice would also be inconsistent with the use of eltype in eg Random, Distributions, and DataStructures.

Maybe it would be best to not define eltype. Or define it as typeof(vi.values) or eltype(vi.values) (defined for NamedTuple and Dict but possibly abstract) for SimpleVarInfo.

A (temporary) workaround for the ::Type issue could be to always base it on typeof(getlogp(vi)) instead of eltype(vi, sampler) in matchingtype (IMO it is a bit unintuitive since it is used for allocating arrays of samples and converting data) or to specialize it for every supported varinfo type (probably introduces inconsistencies).

devmotion avatar Mar 02 '22 10:03 devmotion