Remove Dictionaries with Any key type

Author: penelopeysm

src/DynamicPPL.jl

@@ -23,7 +23,7 @@ using DocStringExtensions
 using Random: Random
 
 # For extending
-import AbstractPPL: predict
+import AbstractPPL: predict, hasvalue, getvalue
 
 # TODO: Remove these when it's possible.
 import Bijectors: link, invlink

src/model.jl

@@ -981,7 +981,11 @@ Base.nameof(model::Model{<:Function}) = nameof(model.f)
 Generate a sample of type `T` from the prior distribution of the `model`.
 """
 function Base.rand(rng::Random.AbstractRNG, ::Type{T}, model::Model) where {T}
-    x = last(evaluate_and_sample!!(rng, model, SimpleVarInfo{Float64}(OrderedDict())))
+    x = last(
+        evaluate_and_sample!!(
+            rng, model, SimpleVarInfo{Float64}(OrderedDict{VarName,Any}())
+        ),
+    )
     return values_as(x, T)
 end
 
@@ -1032,7 +1036,7 @@ julia> logjoint(demo_model([1., 2.]), chain);
 function logjoint(model::Model, chain::AbstractMCMC.AbstractChains)
     var_info = VarInfo(model) # extract variables info from the model
     map(Iterators.product(1:size(chain, 1), 1:size(chain, 3))) do (iteration_idx, chain_idx)
-        argvals_dict = OrderedDict(
+        argvals_dict = OrderedDict{VarName,Any}(
             vn_parent =>
                 values_from_chain(var_info, vn_parent, chain, chain_idx, iteration_idx) for
             vn_parent in keys(var_info)
@@ -1090,7 +1094,7 @@ julia> logprior(demo_model([1., 2.]), chain);
 function logprior(model::Model, chain::AbstractMCMC.AbstractChains)
     var_info = VarInfo(model) # extract variables info from the model
     map(Iterators.product(1:size(chain, 1), 1:size(chain, 3))) do (iteration_idx, chain_idx)
-        argvals_dict = OrderedDict(
+        argvals_dict = OrderedDict{VarName,Any}(
             vn_parent =>
                 values_from_chain(var_info, vn_parent, chain, chain_idx, iteration_idx) for
             vn_parent in keys(var_info)
@@ -1144,7 +1148,7 @@ julia> loglikelihood(demo_model([1., 2.]), chain);
 function Distributions.loglikelihood(model::Model, chain::AbstractMCMC.AbstractChains)
     var_info = VarInfo(model) # extract variables info from the model
     map(Iterators.product(1:size(chain, 1), 1:size(chain, 3))) do (iteration_idx, chain_idx)
-        argvals_dict = OrderedDict(
+        argvals_dict = OrderedDict{VarName,Any}(
             vn_parent =>
                 values_from_chain(var_info, vn_parent, chain, chain_idx, iteration_idx) for
             vn_parent in keys(var_info)

src/simple_varinfo.jl

@@ -62,7 +62,7 @@ ERROR: type NamedTuple has no field x
 [...]
 
 julia> # If one does not know the varnames, we can use a `OrderedDict` instead.
-       _, vi = DynamicPPL.evaluate_and_sample!!(rng, m, SimpleVarInfo{Float64}(OrderedDict()));
+       _, vi = DynamicPPL.evaluate_and_sample!!(rng, m, SimpleVarInfo{Float64}(OrderedDict{VarName,Any}()));
 
 julia> # (✓) Sort of fast, but only possible at runtime.
        vi[@varname(x[1])]
@@ -107,7 +107,7 @@ julia> any(xs .< 0)  # (✓) Positive probability mass on negative numbers!
 true
 
 julia> # And with `OrderedDict` of course!
-       _, vi = DynamicPPL.evaluate_and_sample!!(rng, m, DynamicPPL.settrans!!(SimpleVarInfo(OrderedDict()), true));
+       _, vi = DynamicPPL.evaluate_and_sample!!(rng, m, DynamicPPL.settrans!!(SimpleVarInfo(OrderedDict{VarName,Any}()), true));
 
 julia> vi[@varname(x)] # (✓) -∞ < x < ∞
 0.6225185067787314
@@ -212,7 +212,7 @@ end
 function SimpleVarInfo(values)
     return SimpleVarInfo{LogProbType}(values)
 end
-function SimpleVarInfo(values::Union{<:NamedTuple,<:AbstractDict})
+function SimpleVarInfo(values::Union{<:NamedTuple,<:AbstractDict{<:VarName}})
     return if isempty(values)
         # Can't infer from values, so we just use default.
         SimpleVarInfo{LogProbType}(values)
@@ -264,7 +264,7 @@ function SimpleVarInfo{T}(vi::NTVarInfo, ::Type{D}) where {T<:Real,D}
 end
 
 function untyped_simple_varinfo(model::Model)
-    varinfo = SimpleVarInfo(OrderedDict())
+    varinfo = SimpleVarInfo(OrderedDict{VarName,Any}())
     return last(evaluate_and_sample!!(model, varinfo))
 end
 

src/test_utils/varinfo.jl

@@ -34,7 +34,7 @@ function setup_varinfos(
 
     # SimpleVarInfo
     svi_typed = SimpleVarInfo(example_values)
-    svi_untyped = SimpleVarInfo(OrderedDict())
+    svi_untyped = SimpleVarInfo(OrderedDict{VarName,Any}())
     svi_vnv = SimpleVarInfo(DynamicPPL.VarNamedVector())
 
     varinfos = map((

src/values_as_in_model.jl

@@ -12,12 +12,12 @@ $(TYPEDFIELDS)
 """
 struct ValuesAsInModelAccumulator <: AbstractAccumulator
     "values that are extracted from the model"
-    values::OrderedDict
+    values::OrderedDict{<:VarName}
     "whether to extract variables on the LHS of :="
     include_colon_eq::Bool
 end
 function ValuesAsInModelAccumulator(include_colon_eq)
-    return ValuesAsInModelAccumulator(OrderedDict(), include_colon_eq)
+    return ValuesAsInModelAccumulator(OrderedDict{VarName,Any}(), include_colon_eq)
 end
 
 function Base.copy(acc::ValuesAsInModelAccumulator)

src/varnamedvector.jl

@@ -1482,7 +1482,7 @@ function values_as(vnv::VarNamedVector, ::Type{D}) where {D<:AbstractDict}
 end
 
 # See the docstring of `getvalue` for the semantics of `hasvalue` and `getvalue`, and how
-# they differ from `haskey` and `getindex`. They can be found in src/utils.jl.
+# they differ from `haskey` and `getindex`. They can be found in AbstractPPL.jl.
 
 # TODO(mhauru) This is tricky to implement in the general case, and the below implementation
 # only covers some simple cases. It's probably sufficient in most situations though.

test/runtests.jl

@@ -27,6 +27,9 @@ using LinearAlgebra # Diagonal
 
 using JET: JET
 
+# need to call this to get the AbstractPPL I think
+Pkg.update()
+
 using Combinatorics: combinations
 using OrderedCollections: OrderedSet
 

test/simple_varinfo.jl

@@ -90,7 +90,7 @@
                                                           DynamicPPL.TestUtils.DEMO_MODELS
         values_constrained = DynamicPPL.TestUtils.rand_prior_true(model)
         @testset "$name" for (name, vi) in (
-            ("SVI{Dict}", SimpleVarInfo(Dict())),
+            ("SVI{Dict}", SimpleVarInfo(Dict{VarName,Any}())),
             ("SVI{NamedTuple}", SimpleVarInfo(values_constrained)),
             ("SVI{VNV}", SimpleVarInfo(DynamicPPL.VarNamedVector())),
             ("TypedVarInfo", DynamicPPL.typed_varinfo(model)),

test/varinfo.jl

@@ -110,7 +110,7 @@ end
         test_base(VarInfo())
         test_base(DynamicPPL.typed_varinfo(VarInfo()))
         test_base(SimpleVarInfo())
-        test_base(SimpleVarInfo(Dict()))
+        test_base(SimpleVarInfo(Dict{VarName,Any}()))
         test_base(SimpleVarInfo(DynamicPPL.VarNamedVector()))
     end
 
@@ -597,7 +597,7 @@ end
         test_linked_varinfo(model, vi)
 
         ## `SimpleVarInfo{<:Dict}`
-        vi = DynamicPPL.settrans!!(SimpleVarInfo(Dict()), true)
+        vi = DynamicPPL.settrans!!(SimpleVarInfo(Dict{VarName,Any}()), true)
         test_linked_varinfo(model, vi)
 
         ## `SimpleVarInfo{<:VarNamedVector}`
@@ -737,11 +737,10 @@ end
             model, (; x=1.0), (@varname(x),); include_threadsafe=true
         )
         @testset "$(short_varinfo_name(varinfo))" for varinfo in varinfos
-            # Skip the severely inconcrete `SimpleVarInfo` types, since checking for type
+            # Skip the inconcrete `SimpleVarInfo` types, since checking for type
             # stability for them doesn't make much sense anyway.
-            if varinfo isa SimpleVarInfo{OrderedDict{Any,Any}} ||
-                varinfo isa
-               DynamicPPL.ThreadSafeVarInfo{<:SimpleVarInfo{OrderedDict{Any,Any}}}
+            if varinfo isa SimpleVarInfo{<:AbstractDict} ||
+                varinfo isa DynamicPPL.ThreadSafeVarInfo{<:SimpleVarInfo{<:AbstractDict}}
                 continue
             end
             @inferred DynamicPPL.unflatten(varinfo, varinfo[:])