renamed getval to getindex_internal and made dist an optional

argument for all the transform-related methods

Author: torfjelde

src/abstract_varinfo.jl

@@ -779,14 +779,14 @@ end
 Invlink `x` and compute the logpdf under `dist` including correction from
 the invlink-transformation.
 
-If `x` is not provided, `getval(vi, vn)` will be used.
+If `x` is not provided, `getindex_internal(vi, vn)` will be used.
 
 !!! warning
     The input value `x` should be according to the internal representation of
-    `varinfo`, e.g. the value returned by `getval(vi, vn)`.
+    `varinfo`, e.g. the value returned by `getindex_internal(vi, vn)`.
 """
 function invlink_with_logpdf(vi::AbstractVarInfo, vn::VarName, dist)
-    return invlink_with_logpdf(vi, vn, dist, getval(vi, vn))
+    return invlink_with_logpdf(vi, vn, dist, getindex_internal(vi, vn))
 end
 function invlink_with_logpdf(vi::AbstractVarInfo, vn::VarName, dist, y)
     f = from_maybe_linked_internal_transform(vi, vn, dist)
@@ -800,26 +800,32 @@ increment_num_produce!(::AbstractVarInfo) = nothing
 setgid!(vi::AbstractVarInfo, gid::Selector, vn::VarName) = nothing
 
 """
-    from_internal_transform(varinfo::AbstractVarInfo, vn::VarName, dist)
+    from_internal_transform(varinfo::AbstractVarInfo, vn::VarName[, dist])
 
 Return a transformation that transforms from the internal representation of `vn` with `dist`
 in `varinfo` to a representation compatible with `dist`.
+
+If `dist` is not present, then it is assumed that `varinfo` knows the correct output for `vn`.
 """
 function from_internal_transform end
 
 """
-    from_linked_internal_transform(varinfo::AbstractVarInfo, vn::VarName, dist)
+    from_linked_internal_transform(varinfo::AbstractVarInfo, vn::VarName[, dist])
 
 Return a transformation that transforms from the linked internal representation of `vn` with `dist`
 in `varinfo` to a representation compatible with `dist`.
+
+If `dist` is not present, then it is assumed that `varinfo` knows the correct output for `vn`.
 """
 function from_linked_internal_transform end
 
 """
-    from_maybe_linked_internal_transform(varinfo::AbstractVarInfo, vn::VarName, dist)
+    from_maybe_linked_internal_transform(varinfo::AbstractVarInfo, vn::VarName[, dist])
 
 Return a transformation that transforms from the possibly linked internal representation of `vn` with `dist`n
 in `varinfo` to a representation compatible with `dist`.
+
+If `dist` is not present, then it is assumed that `varinfo` knows the correct output for `vn`.
 """
 function from_maybe_linked_internal_transform(varinfo::AbstractVarInfo, vn::VarName, dist)
     return if istrans(varinfo, vn)
@@ -828,57 +834,94 @@ function from_maybe_linked_internal_transform(varinfo::AbstractVarInfo, vn::VarN
         from_internal_transform(varinfo, vn, dist)
     end
 end
+function from_maybe_linked_internal_transform(varinfo::AbstractVarInfo, vn::VarName)
+    return if istrans(varinfo, vn)
+        from_linked_internal_transform(varinfo, vn)
+    else
+        from_internal_transform(varinfo, vn)
+    end
+end
 
 """
-    to_internal_transform(varinfo::AbstractVarInfo, vn::VarName, dist)
+    to_internal_transform(varinfo::AbstractVarInfo, vn::VarName[, dist])
 
 Return a transformation that transforms from a representation compatible with `dist` to the
 internal representation of `vn` with `dist` in `varinfo`.
+
+If `dist` is not present, then it is assumed that `varinfo` knows the correct output for `vn`.
 """
 function to_internal_transform(varinfo::AbstractVarInfo, vn::VarName, dist)
     return inverse(from_internal_transform(varinfo, vn, dist))
 end
+function to_internal_transform(varinfo::AbstractVarInfo, vn::VarName)
+    return inverse(from_internal_transform(varinfo, vn))
+end
 
 """
-    to_linked_internal_transform(varinfo::AbstractVarInfo, vn::VarName, dist)
+    to_linked_internal_transform(varinfo::AbstractVarInfo, vn::VarName[, dist])
 
 Return a transformation that transforms from a representation compatible with `dist` to the
 linked internal representation of `vn` with `dist` in `varinfo`.
+
+If `dist` is not present, then it is assumed that `varinfo` knows the correct output for `vn`.
 """
 function to_linked_internal_transform(varinfo::AbstractVarInfo, vn::VarName, dist)
     return inverse(from_linked_internal_transform(varinfo, vn, dist))
 end
+function to_linked_internal_transform(varinfo::AbstractVarInfo, vn::VarName)
+    return inverse(from_linked_internal_transform(varinfo, vn))
+end
 
 """
-    to_maybe_linked_internal_transform(varinfo::AbstractVarInfo, vn::VarName, dist)
+    to_maybe_linked_internal_transform(varinfo::AbstractVarInfo, vn::VarName[, dist])
 
 Return a transformation that transforms from a representation compatible with `dist` to a
 possibly linked internal representation of `vn` with `dist` in `varinfo`.
+
+If `dist` is not present, then it is assumed that `varinfo` knows the correct output for `vn`.
 """
 function to_maybe_linked_internal_transform(varinfo::AbstractVarInfo, vn::VarName, dist)
     return inverse(from_maybe_linked_internal_transform(varinfo, vn, dist))
 end
+function to_maybe_linked_internal_transform(varinfo::AbstractVarInfo, vn::VarName)
+    return inverse(from_maybe_linked_internal_transform(varinfo, vn))
+end
 
 """
     internal_to_linked_internal_transform(varinfo::AbstractVarInfo, vn::VarName, dist)
 
 Return a transformation that transforms from the internal representation of `vn` with `dist`
 in `varinfo` to a _linked_ internal representation of `vn` with `dist` in `varinfo`.
+
+If `dist` is not present, then it is assumed that `varinfo` knows the correct output for `vn`.
 """
 function internal_to_linked_internal_transform(varinfo::AbstractVarInfo, vn::VarName, dist)
     f_from_internal = from_internal_transform(varinfo, vn, dist)
     f_to_linked_internal = to_linked_internal_transform(varinfo, vn, dist)
     return f_to_linked_internal ∘ f_from_internal
 end
+function internal_to_linked_internal_transform(varinfo::AbstractVarInfo, vn::VarName)
+    f_from_internal = from_internal_transform(varinfo, vn)
+    f_to_linked_internal = to_linked_internal_transform(varinfo, vn)
+    return f_to_linked_internal ∘ f_from_internal
+end
 
 """
-    linked_internal_to_internal_transform(varinfo::AbstractVarInfo, vn::VarName, dist)
+    linked_internal_to_internal_transform(varinfo::AbstractVarInfo, vn::VarName[, dist])
 
 Return a transformation that transforms from a _linked_ internal representation of `vn` with `dist`
 in `varinfo` to the internal representation of `vn` with `dist` in `varinfo`.
+
+If `dist` is not present, then it is assumed that `varinfo` knows the correct output for `vn`.
 """
 function linked_internal_to_internal_transform(varinfo::AbstractVarInfo, vn::VarName, dist)
     f_from_linked_internal = from_linked_internal_transform(varinfo, vn, dist)
     f_to_internal = to_internal_transform(varinfo, vn, dist)
     return f_to_internal ∘ f_from_linked_internal
 end
+
+function linked_internal_to_internal_transform(varinfo::AbstractVarInfo, vn::VarName)
+    f_from_linked_internal = from_linked_internal_transform(varinfo, vn)
+    f_to_internal = to_internal_transform(varinfo, vn)
+    return f_to_internal ∘ f_from_linked_internal
+end

src/context_implementations.jl

@@ -477,7 +477,7 @@ function _link_broadcast_new(vi, vn, dist, r)
 end
 
 function _maybe_invlink_broadcast(vi, vn, dist)
-    xvec = getval(vi, vn)
+    xvec = getindex_internal(vi, vn)
     b = from_maybe_linked_internal_transform(vi, vn, dist)
     return b(xvec)
 end

src/simple_varinfo.jl

@@ -332,7 +332,7 @@ function getindex_raw(vi::SimpleVarInfo, vns::Vector{<:VarName}, dist::Distribut
 end
 
 # HACK: because `VarInfo` isn't ready to implement a proper `getindex_raw`.
-getval(vi::SimpleVarInfo, vn::VarName) = getindex_raw(vi, vn)
+getindex_internal(vi::SimpleVarInfo, vn::VarName) = getindex_raw(vi, vn)
 
 Base.haskey(vi::SimpleVarInfo, vn::VarName) = hasvalue(vi.values, vn)
 
@@ -690,16 +690,11 @@ function invlink!!(
 end
 
 # With `SimpleVarInfo`, when we're not working with linked variables, there's no need to do anything.
-from_internal_transform(::SimpleVarInfo, dist) = identity
-function from_internal_transform(vi::SimpleVarInfo, ::VarName, dist)
-    return from_internal_transform(vi, dist)
-end
-
-function from_linked_internal_transform(vi::SimpleVarInfo, dist)
-    return invlink_transform(dist)
-end
+from_internal_transform(vi::SimpleVarInfo, ::VarName) = identity
+from_internal_transform(vi::SimpleVarInfo, ::VarName, dist) = identity
+from_linked_internal_transform(vi::SimpleVarInfo, ::VarName) = identity
 function from_linked_internal_transform(vi::SimpleVarInfo, ::VarName, dist)
-    return from_linked_internal_transform(vi, dist)
+    return invlink_transform(dist)
 end
 
 # Threadsafe stuff.

src/threadsafe.jl

@@ -222,7 +222,7 @@ end
 istrans(vi::ThreadSafeVarInfo, vn::VarName) = istrans(vi.varinfo, vn)
 istrans(vi::ThreadSafeVarInfo, vns::AbstractVector{<:VarName}) = istrans(vi.varinfo, vns)
 
-getval(vi::ThreadSafeVarInfo, vn::VarName) = getval(vi.varinfo, vn)
+getindex_internal(vi::ThreadSafeVarInfo, vn::VarName) = getindex_internal(vi.varinfo, vn)
 
 function unflatten(vi::ThreadSafeVarInfo, x::AbstractVector)
     return Setfield.@set vi.varinfo = unflatten(vi.varinfo, x)

src/utils.jl

@@ -256,10 +256,14 @@ from_vec_transform(dist::Distribution) = from_vec_transform_for_size(size(dist))
 from_vec_transform(dist::LKJCholesky) = ToChol(dist.uplo) ∘ FromVec(size(dist))
 
 """
-    from_linked_vec_transform(dist)
+    from_linked_vec_transform(dist::Distribution)
 
 Return the transformation from the unconstrained vector to the constrained
 realization of distribution `dist`.
+
+By default, this is just `invlink_transform(dist) ∘ from_vec_transform(dist)`.
+
+See also: [`DynamicPPL.invlink_transform`](@ref), [`DynamicPPL.from_vec_transform`](@ref).
 """
 function from_linked_vec_transform(dist::Distribution)
     f_vec = from_vec_transform(dist)