Remove selector stuff from VarInfo tests and link/invlink

docs/src/api.md

@@ -304,7 +304,7 @@ set_num_produce!
 increment_num_produce!
 reset_num_produce!
 setorder!
-set_retained_vns_del_by_spl!
+set_retained_vns_del!
 ```
 
 ```@docs

src/DynamicPPL.jl

@@ -59,7 +59,7 @@ export AbstractVarInfo,
     set_num_produce!,
     reset_num_produce!,
     increment_num_produce!,
-    set_retained_vns_del_by_spl!,
+    set_retained_vns_del!,
     is_flagged,
     set_flag!,
     unset_flag!,

src/abstract_varinfo.jl

@@ -537,113 +537,116 @@
 """
 function settrans!! end
 
+# For link!!, invlink!!, link, and invlink, we deliberately do not provide a fallback
+# method for the case when no `vns` is provided, that would get all the keys from the
+# `VarInfo`. Hence each subtype of `AbstractVarInfo` needs to implement separately the case
+# where `vns` is provided and the one where it is not. This is because having separate
+# implementations is typically much more performant, and because not all AbstractVarInfo
+# types support partial linking.
+
 """
     link!!([t::AbstractTransformation, ]vi::AbstractVarInfo, model::Model)
-    link!!([t::AbstractTransformation, ]vi::AbstractVarInfo, spl::AbstractSampler, model::Model)
+    link!!([t::AbstractTransformation, ]vi::AbstractVarInfo, model::Model)
+    link!!([t::AbstractTransformation, ]vi::AbstractVarInfo, vns::Tuple{N,VarName}, model::Model)
 
-Transform the variables in `vi` to their linked space, using the transformation `t`,
-mutating `vi` if possible.
+Transform variables in `vi` to their linked space, mutating `vi` if possible.
 
-If `t` is not provided, `default_transformation(model, vi)` will be used.
+Either transform all variables, or only ones specified in `vns`.
+
+Use the  transformation `t`, or `default_transformation(model, vi)` if one is not provided.
 
 See also: [`default_transformation`](@ref), [`invlink!!`](@ref).
 """
-link!!(vi::AbstractVarInfo, model::Model) = link!!(vi, SampleFromPrior(), model)
-function link!!(t::AbstractTransformation, vi::AbstractVarInfo, model::Model)
-    return link!!(t, vi, SampleFromPrior(), model)
+function link!!(vi::AbstractVarInfo, model::Model)
+    return link!!(default_transformation(model, vi), vi, model)
 end
-function link!!(vi::AbstractVarInfo, spl::AbstractSampler, model::Model)
-    # Use `default_transformation` to decide which transformation to use if none is specified.
-    return link!!(default_transformation(model, vi), vi, spl, model)
+function link!!(vi::AbstractVarInfo, vns::VarNameTuple, model::Model)
+    return link!!(default_transformation(model, vi), vi, vns, model)
 end
 
 """
     link([t::AbstractTransformation, ]vi::AbstractVarInfo, model::Model)
-    link([t::AbstractTransformation, ]vi::AbstractVarInfo, spl::AbstractSampler, model::Model)
+    link([t::AbstractTransformation, ]vi::AbstractVarInfo, model::Model)
+    link([t::AbstractTransformation, ]vi::AbstractVarInfo, vns::Tuple{N,VarName}, model::Model)
 
-Transform the variables in `vi` to their linked space without mutating `vi`, using the transformation `t`.
+Transform variables in `vi` to their linked space without mutating `vi`.
 
-If `t` is not provided, `default_transformation(model, vi)` will be used.
+Either transform all variables, or only ones specified in `vns`.
+
+Use the  transformation `t`, or `default_transformation(model, vi)` if one is not provided.
 
 See also: [`default_transformation`](@ref), [`invlink`](@ref).
 """
-link(vi::AbstractVarInfo, model::Model) = link(vi, SampleFromPrior(), model)
-function link(t::AbstractTransformation, vi::AbstractVarInfo, model::Model)
-    return link(t, deepcopy(vi), SampleFromPrior(), model)
+function link(vi::AbstractVarInfo, model::Model)
+    return link(default_transformation(model, vi), vi, model)
 end
-function link(vi::AbstractVarInfo, spl::AbstractSampler, model::Model)
-    # Use `default_transformation` to decide which transformation to use if none is specified.
-    return link(default_transformation(model, vi), deepcopy(vi), spl, model)
+function link(vi::AbstractVarInfo, vns::VarNameTuple, model::Model)
+    return link(default_transformation(model, vi), vi, vns, model)
 end
 
 """
     invlink!!([t::AbstractTransformation, ]vi::AbstractVarInfo, model::Model)
-    invlink!!([t::AbstractTransformation, ]vi::AbstractVarInfo, spl::AbstractSampler, model::Model)
+    invlink!!([t::AbstractTransformation, ]vi::AbstractVarInfo, vns::Tuple{N,VarName}, model::Model)
+
+Transform variables in `vi` to their constrained space, mutating `vi` if possible.
 
-Transform the variables in `vi` to their constrained space, using the (inverse of)
-transformation `t`, mutating `vi` if possible.
+Either transform all variables, or only ones specified in `vns`.
 
-If `t` is not provided, `default_transformation(model, vi)` will be used.
+Use the (inverse of) transformation `t`, or `default_transformation(model, vi)` if one is
+not provided.
 
 See also: [`default_transformation`](@ref), [`link!!`](@ref).
 """
-invlink!!(vi::AbstractVarInfo, model::Model) = invlink!!(vi, SampleFromPrior(), model)
-function invlink!!(t::AbstractTransformation, vi::AbstractVarInfo, model::Model)
-    return invlink!!(t, vi, SampleFromPrior(), model)
+function invlink!!(vi::AbstractVarInfo, model::Model)
+    return invlink!!(default_transformation(model, vi), vi, model)
 end
-function invlink!!(vi::AbstractVarInfo, spl::AbstractSampler, model::Model)
-    # Here we extract the `transformation` from `vi` rather than using the default one.
-    return invlink!!(transformation(vi), vi, spl, model)
+function invlink!!(vi::AbstractVarInfo, vns::VarNameTuple, model::Model)
+    return invlink!!(default_transformation(model, vi), vi, vns, model)
 end
 
 # Vector-based ones.
 function link!!(
-    t::StaticTransformation{<:Bijectors.Transform},
-    vi::AbstractVarInfo,
-    spl::AbstractSampler,
-    model::Model,
+    t::StaticTransformation{<:Bijectors.Transform}, vi::AbstractVarInfo, ::Model
 )
     b = inverse(t.bijector)
-    x = vi[spl]
+    x = vi[:]
     y, logjac = with_logabsdet_jacobian(b, x)
 
     lp_new = getlogp(vi) - logjac
-    vi_new = setlogp!!(unflatten(vi, spl, y), lp_new)
+    vi_new = setlogp!!(unflatten(vi, y), lp_new)
     return settrans!!(vi_new, t)
 end
 
 function invlink!!(
-    t::StaticTransformation{<:Bijectors.Transform},
-    vi::AbstractVarInfo,
-    spl::AbstractSampler,
-    model::Model,
+    t::StaticTransformation{<:Bijectors.Transform}, vi::AbstractVarInfo, ::Model
 )
     b = t.bijector
-    y = vi[spl]
+    y = vi[:]
     x, logjac = with_logabsdet_jacobian(b, y)
 
     lp_new = getlogp(vi) + logjac
-    vi_new = setlogp!!(unflatten(vi, spl, x), lp_new)
+    vi_new = setlogp!!(unflatten(vi, x), lp_new)
     return settrans!!(vi_new, NoTransformation())
 end
 
 """
     invlink([t::AbstractTransformation, ]vi::AbstractVarInfo, model::Model)
-    invlink([t::AbstractTransformation, ]vi::AbstractVarInfo, spl::AbstractSampler, model::Model)
+    invlink([t::AbstractTransformation, ]vi::AbstractVarInfo, vns::Tuple{N,VarName}, model::Model)
+
+Transform variables in `vi` to their constrained space without mutating `vi`.
 
-Transform the variables in `vi` to their constrained space without mutating `vi`, using the (inverse of)
-transformation `t`.
+Either transform all variables, or only ones specified in `vns`.
 
-If `t` is not provided, `default_transformation(model, vi)` will be used.
+Use the (inverse of) transformation `t`, or `default_transformation(model, vi)` if one is
+not provided.
 
 See also: [`default_transformation`](@ref), [`link`](@ref).
 """
-invlink(vi::AbstractVarInfo, model::Model) = invlink(vi, SampleFromPrior(), model)
-function invlink(t::AbstractTransformation, vi::AbstractVarInfo, model::Model)
-    return invlink(t, vi, SampleFromPrior(), model)
+function invlink(vi::AbstractVarInfo, model::Model)
+    return invlink(default_transformation(model, vi), vi, model)
 end
-function invlink(vi::AbstractVarInfo, spl::AbstractSampler, model::Model)
-    return invlink(transformation(vi), vi, spl, model)
+function invlink(vi::AbstractVarInfo, vns::VarNameTuple, model::Model)
+    return invlink(default_transformation(model, vi), vi, vns, model)
 end
 
 """
@@ -715,9 +718,10 @@
     return vi
 end
 function maybe_invlink_before_eval!!(
-    t::StaticTransformation, vi::AbstractVarInfo, context::AbstractContext, model::Model
+    t::StaticTransformation, vi::AbstractVarInfo, ::AbstractContext, model::Model
 )
-    return invlink!!(t, vi, _default_sampler(context), model)
+    # TODO(mhauru) Why does this function need the context argument?
+    return invlink!!(t, vi, model)
 end
 
 function _default_sampler(context::AbstractContext)

src/simple_varinfo.jl

@@ -680,8 +680,7 @@ Distributions.loglikelihood(model::Model, θ) = loglikelihood(model, SimpleVarIn
 function link!!(
     t::StaticTransformation{<:Bijectors.NamedTransform},
     vi::SimpleVarInfo{<:NamedTuple},
-    spl::AbstractSampler,
-    model::Model,
+    ::Model,
 )
     # TODO: Make sure that `spl` is respected.
     b = inverse(t.bijector)
@@ -695,8 +694,7 @@ end
 function invlink!!(
     t::StaticTransformation{<:Bijectors.NamedTransform},
     vi::SimpleVarInfo{<:NamedTuple},
-    spl::AbstractSampler,
-    model::Model,
+    ::Model,
 )
     # TODO: Make sure that `spl` is respected.
     b = t.bijector

src/threadsafe.jl

@@ -81,59 +81,43 @@
 
 islinked(vi::ThreadSafeVarInfo, spl::AbstractSampler) = islinked(vi.varinfo, spl)
 
-function link!!(
-    t::AbstractTransformation, vi::ThreadSafeVarInfo, spl::AbstractSampler, model::Model
-)
-    return Accessors.@set vi.varinfo = link!!(t, vi.varinfo, spl, model)
+function link!!(t::AbstractTransformation, vi::ThreadSafeVarInfo, args...)
+    return Accessors.@set vi.varinfo = link!!(t, vi.varinfo, args...)
 end
 
-function invlink!!(
-    t::AbstractTransformation, vi::ThreadSafeVarInfo, spl::AbstractSampler, model::Model
-)
-    return Accessors.@set vi.varinfo = invlink!!(t, vi.varinfo, spl, model)
+function invlink!!(t::AbstractTransformation, vi::ThreadSafeVarInfo, args...)
+    return Accessors.@set vi.varinfo = invlink!!(t, vi.varinfo, args...)
 end
 
-function link(
-    t::AbstractTransformation, vi::ThreadSafeVarInfo, spl::AbstractSampler, model::Model
-)
-    return Accessors.@set vi.varinfo = link(t, vi.varinfo, spl, model)
+function link(t::AbstractTransformation, vi::ThreadSafeVarInfo, args...)
+    return Accessors.@set vi.varinfo = link(t, vi.varinfo, args...)
 end
 
-function invlink(
-    t::AbstractTransformation, vi::ThreadSafeVarInfo, spl::AbstractSampler, model::Model
-)
-    return Accessors.@set vi.varinfo = invlink(t, vi.varinfo, spl, model)
+function invlink(t::AbstractTransformation, vi::ThreadSafeVarInfo, args...)
+    return Accessors.@set vi.varinfo = invlink(t, vi.varinfo, args...)
 end
 
 # Need to define explicitly for `DynamicTransformation` to avoid method ambiguity.
 # NOTE: We also can't just defer to the wrapped varinfo, because we need to ensure
 # consistency between `vi.logps` field and `getlogp(vi.varinfo)`, which accumulates
 # to define `getlogp(vi)`.
-function link!!(
-    t::DynamicTransformation, vi::ThreadSafeVarInfo, spl::AbstractSampler, model::Model
-)
+function link!!(t::DynamicTransformation, vi::ThreadSafeVarInfo, model::Model)
     return settrans!!(last(evaluate!!(model, vi, DynamicTransformationContext{false}())), t)
 end
 
-function invlink!!(
-    ::DynamicTransformation, vi::ThreadSafeVarInfo, spl::AbstractSampler, model::Model
-)
+function invlink!!(::DynamicTransformation, vi::ThreadSafeVarInfo, model::Model)
     return settrans!!(
         last(evaluate!!(model, vi, DynamicTransformationContext{true}())),
         NoTransformation(),
     )
 end
 
-function link(
-    t::DynamicTransformation, vi::ThreadSafeVarInfo, spl::AbstractSampler, model::Model
-)
-    return link!!(t, deepcopy(vi), spl, model)
+function link(t::DynamicTransformation, vi::ThreadSafeVarInfo, model::Model)
+    return link!!(t, deepcopy(vi), model)
 end
 
-function invlink(
-    t::DynamicTransformation, vi::ThreadSafeVarInfo, spl::AbstractSampler, model::Model
-)
-    return invlink!!(t, deepcopy(vi), spl, model)
+function invlink(t::DynamicTransformation, vi::ThreadSafeVarInfo, model::Model)
+    return invlink!!(t, deepcopy(vi), model)
 end
 
 function maybe_invlink_before_eval!!(
@@ -182,8 +166,8 @@
     return vector_getranges(vi.varinfo, vns)
 end
 
-function set_retained_vns_del_by_spl!(vi::ThreadSafeVarInfo, spl::Sampler)
-    return set_retained_vns_del_by_spl!(vi.varinfo, spl)
+function set_retained_vns_del!(vi::ThreadSafeVarInfo, spl::Sampler)
+    return set_retained_vns_del!(vi.varinfo, spl)
 end
 
 isempty(vi::ThreadSafeVarInfo) = isempty(vi.varinfo)

src/transforming.jl

@@ -91,29 +91,21 @@ function dot_tilde_assume(
     return r, lp, vi
 end
 
-function link!!(
-    t::DynamicTransformation, vi::AbstractVarInfo, spl::AbstractSampler, model::Model
-)
+function link!!(t::DynamicTransformation, vi::AbstractVarInfo, model::Model)
     return settrans!!(last(evaluate!!(model, vi, DynamicTransformationContext{false}())), t)
 end
 
-function invlink!!(
-    ::DynamicTransformation, vi::AbstractVarInfo, spl::AbstractSampler, model::Model
-)
+function invlink!!(::DynamicTransformation, vi::AbstractVarInfo, model::Model)
     return settrans!!(
         last(evaluate!!(model, vi, DynamicTransformationContext{true}())),
         NoTransformation(),
     )
 end
 
-function link(
-    t::DynamicTransformation, vi::AbstractVarInfo, spl::AbstractSampler, model::Model
-)
-    return link!!(t, deepcopy(vi), spl, model)
+function link(t::DynamicTransformation, vi::AbstractVarInfo, model::Model)
+    return link!!(t, deepcopy(vi), model)
 end
 
-function invlink(
-    t::DynamicTransformation, vi::AbstractVarInfo, spl::AbstractSampler, model::Model
-)
-    return invlink!!(t, deepcopy(vi), spl, model)
+function invlink(t::DynamicTransformation, vi::AbstractVarInfo, model::Model)
+    return invlink!!(t, deepcopy(vi), model)
 end