Fix docs, workaround Zygote issue

Author: mhauru

docs/src/internals/varinfo.md

@@ -85,7 +85,9 @@ For example, with the model above we have
 
 ```@example varinfo-design
 # Type-unstable `VarInfo`
-varinfo_untyped = DynamicPPL.untyped_varinfo(demo())
+varinfo_untyped = DynamicPPL.untyped_varinfo(
+    demo(), SampleFromPrior(), DefaultContext(), DynamicPPL.Metadata
+)
 typeof(varinfo_untyped.metadata)
 ```
 

src/utils.jl

@@ -297,8 +297,8 @@ from_vec_transform_for_size(::Tuple{<:Any}) = identity
 Return the transformation from the vector representation of a realization from
 distribution `dist` to the original representation compatible with `dist`.
 """
-from_vec_transform(::UnivariateDistribution) = UnwrapSingletonTransform()
 from_vec_transform(dist::Distribution) = from_vec_transform_for_size(size(dist))
+from_vec_transform(::UnivariateDistribution) = UnwrapSingletonTransform()
 from_vec_transform(dist::LKJCholesky) = ToChol(dist.uplo) ∘ ReshapeTransform(size(dist))
 
 """

src/varinfo.jl

@@ -184,7 +184,7 @@ function untyped_varinfo(
     varinfo = VarInfo(metadata_type())
     return last(evaluate!!(model, varinfo, SamplingContext(rng, sampler, context)))
 end
-function untyped_varinfo(model::Model, args::Union{AbstractSampler,AbstractContext}...)
+function untyped_varinfo(model::Model, args::Union{AbstractSampler,AbstractContext,Type}...)
     return untyped_varinfo(Random.default_rng(), model, args...)
 end