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
 
@@ -1028,7 +1032,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)
@@ -1082,7 +1086,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)
@@ -1136,7 +1140,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
@@ -206,7 +206,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)
@@ -258,7 +258,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
 
 accumulator_name(::Type{<:ValuesAsInModelAccumulator}) = :ValuesAsInModel

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/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
 
@@ -604,8 +604,7 @@ end
 
         ## `SimpleVarInfo{<:Dict}`
         vi = DynamicPPL.settrans!!(SimpleVarInfo(Dict()), true)
-        # Sample in unconstrained space.
-        vi = last(DynamicPPL.evaluate_and_sample!!(model, vi))
+        vi = DynamicPPL.settrans!!(SimpleVarInfo(Dict{VarName,Any}()), true)
         f = DynamicPPL.from_linked_internal_transform(vi, vn, dist)
         x = f(DynamicPPL.getindex_internal(vi, vn))
         @test getlogjoint(vi) ≈ Bijectors.logpdf_with_trans(dist, x, true)
@@ -750,11 +749,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[:])