Un-fix predict on varinfo

Author: penelopeysm

src/model.jl

@@ -1200,15 +1200,24 @@ function predict(
     varinfo = DynamicPPL.VarInfo(model)
     return map(chain) do params_varinfo
         vi = deepcopy(varinfo)
-        # TODO(penelopeysm): Requires two model evaluations, one to extract the
-        # parameters and one to set them. The reason why we need values_as_in_model
-        # is because `params_varinfo` may well have some weird combination of
-        # linked/unlinked, whereas `varinfo` is always unlinked since it is
-        # freshly constructed.
-        # This is quite inefficient. It would of course be alright if
-        # ValuesAsInModelAccumulator was a default acc.
-        values_nt = values_as_in_model(model, false, params_varinfo)
-        _, vi = DynamicPPL.init!!(rng, model, vi, ParamsInit(values_nt, PriorInit()))
+        # TODO(penelopeysm): BEWARE - Because `ParamsInit` expects unlinked
+        # values, this is bugged in the case where `params_varinfo` is linked.
+        # This could be solved by using `values_as_in_model`. However, we run
+        # into another problem:
+        #  - The `model` passed into this function will have the target
+        #    prediction variables (say `y`) set to `missing`.
+        #  - Calling `values_as_in_model` on this new model will lead to
+        #    DynamicPPL attempting to extract the values of `y` from the
+        #    `params_varinfo`, which will fail since `y` will not have been
+        #    present.
+        #  I think that the solution is to pass in the ORIGINAL model, but ALSO
+        #  pass in a set / iterable of VarNames that are to be predicted against.
+        #  That way, we can:
+        #  - Safely use `values_as_in_model` to extract a dict of values.
+        #  - Drop the values of the prediction variables from the dict.
+        #  - Then use `ParamsInit` to generate the predictions.
+        values_dict = values_as(params_varinfo, Dict{VarName,Any})
+        _, vi = DynamicPPL.init!!(rng, model, vi, ParamsInit(values_dict, PriorInit()))
         return vi
     end
 end