added a bunch of docs for introduced and existing methods

Added docs for `determine_suitable_varinfo` and existing methods that should be
documented, e.g. `untyped_varinfo`, `typed_varinfo`, and `default_varinfo`

Author: torfjelde

docs/src/api.md

@@ -243,6 +243,13 @@ AbstractVarInfo
 
 But exactly how a [`AbstractVarInfo`](@ref) stores this information can vary.
 
+For constructing the "default" typed and untyped varinfo types used in DynamicPPL (see [the section on varinfo design](./internals/varinfo) for more on this), we have the following two methods:
+
+```@docs
+DynamicPPL.untyped_varinfo
+DynamicPPL.typed_varinfo
+```
+
 #### `VarInfo`
 
 ```@docs
@@ -403,6 +410,19 @@ DynamicPPL.loadstate
 DynamicPPL.initialsampler
 ```
 
+Finally, to specify which varinfo type a [`Sampler`](@ref) should use for a given [`Model`](@ref), this is specified by [`DynamicPPL.default_varinfo`](@ref) and can thus be overloaded for each  `model`-`sampler` combination. This can be useful in cases where one has explicit knowledge that one type of varinfo will be more performant for the given `model` and `sampler`.
+
+```@docs
+DynamicPPL.default_varinfo
+```
+
+There is also the _experimental_ [`DynamicPPL.Experimental.determine_suitable_varinfo`](@ref), which uses static checking via [JET.jl](https://github.com/aviatesk/JET.jl) to determine whether one should use [`DynamicPPL.typed_varinfo`](@ref) or [`DynamicPPL.untyped_varinfo`](@ref), depending on which supports the model:
+
+```@docs
+DynamicPPL.Experimental.determine_suitable_varinfo
+DynamicPPL.Experimental.is_suitable_varinfo
+```
+
 ### [Model-Internal Functions](@id model_internal)
 
 ```@docs

src/sampler.jl

@@ -67,6 +67,20 @@ function AbstractMCMC.step(
     return vi, nothing
 end
 
+"""
+    default_varinfo(rng, model, sampler[, context])
+
+Return a default varinfo object for the given `model` and `sampler`.
+
+# Arguments
+- `rng::Random.AbstractRNG`: Random number generator.
+- `model::Model`: Model for which we want to create a varinfo object.
+- `sampler::AbstractSampler`: Sampler which will make use of the varinfo object.
+- `context::AbstractContext`: Context in which the model is evaluated.
+
+# Returns
+- `AbstractVarInfo`: Default varinfo object for the given `model` and `sampler`.
+"""
 function default_varinfo(rng::Random.AbstractRNG, model::Model, sampler::AbstractSampler)
     return default_varinfo(rng, model, sampler, DefaultContext())
 end

src/varinfo.jl

@@ -163,11 +163,7 @@ function has_varnamedvector(vi::VarInfo)
            (vi isa TypedVarInfo && any(Base.Fix2(isa, VarNamedVector), values(vi.metadata)))
 end
 
-"""
-    untyped_varinfo([rng, ]model[, sampler, context])
-
-Return an untyped `VarInfo` instance for the model `model`.
-"""
+# TODO: We should just remove this function in favour of passing `SamplingContext` directly.
 function untyped_varinfo(
     rng::Random.AbstractRNG,
     model::Model,
@@ -182,6 +178,19 @@ function untyped_varinfo(
 )
     return untyped_varinfo(Random.default_rng(), model, args...)
 end
+
+"""
+    untyped_varinfo([rng, ]model[, context, metadata])
+
+Return an untyped varinfo object for the given `model` and `context`.
+
+# Arguments
+- `rng::Random.AbstractRNG`: The random number generator to use. Default: `Random.default_rng()`.
+- `model::Model`: The model for which to create the varinfo object.
+- `context::AbstractContext`: The context in which to evaluate the model. Default: `DefaultContext()`.
+- `metadata::Union{Metadata,VarNamedVector}`: The metadata to use for the varinfo object.
+    Default: `Metadata()`.
+"""
 function untyped_varinfo(
     model::Model,
     context::AbstractContext,
@@ -194,9 +203,14 @@ function untyped_varinfo(
 end
 
 """
-    typed_varinfo([rng, ]model[, sampler, context])
+    typed_varinfo([rng, ]model[, context, metadata])
+
+Return a typed varinfo object for the given `model`, `sampler` and `context`.
+
+This simply calls [`DynamicPPL.untyped_varinfo`](@ref) and converts the resulting
+varinfo object to a typed varinfo object.
 
-Return a typed `VarInfo` instance for the model `model`.
+See also: [`DynamicPPL.untyped_varinfo`](@ref)
 """
 typed_varinfo(args...) = TypedVarInfo(untyped_varinfo(args...))