Fix performance of getparams on untyped vi (#2605)

Author: penelopeysm

HISTORY.md

@@ -1,3 +1,8 @@
+# Release 0.39.3
+
+Improved the performance of `Turing.Inference.getparams` when called with an untyped VarInfo as the second argument, by first converting to a typed VarInfo.
+This makes, for example, the post-sampling Chains construction for `Prior()` run much faster.
+
 # Release 0.39.2
 
 Fixed a bug in the support of `OrderedLogistic` (by changing the minimum from 0 to 1).

Project.toml

@@ -1,6 +1,6 @@
 name = "Turing"
 uuid = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
-version = "0.39.2"
+version = "0.39.3"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

src/mcmc/Inference.jl

@@ -184,6 +184,14 @@ function getparams(model::DynamicPPL.Model, vi::DynamicPPL.VarInfo)
     # Materialize the iterators and concatenate.
     return mapreduce(collect, vcat, iters)
 end
+function getparams(
+    model::DynamicPPL.Model, untyped_vi::DynamicPPL.VarInfo{<:DynamicPPL.Metadata}
+)
+    # values_as_in_model is unconscionably slow for untyped VarInfo. It's
+    # much faster to convert it to a typed varinfo before calling getparams.
+    # https://github.com/TuringLang/Turing.jl/issues/2604
+    return getparams(model, DynamicPPL.typed_varinfo(untyped_vi))
+end
 function getparams(::DynamicPPL.Model, ::DynamicPPL.VarInfo{NamedTuple{(),Tuple{}}})
     return float(Real)[]
 end