fix definition order

Author: penelopeysm

src/contexts/init.jl

@@ -1,3 +1,71 @@
+"""
+    RangeAndLinked
+
+Suppose we have vectorised parameters `params::AbstractVector{<:Real}`. Each random variable
+in the model will in general correspond to a sub-vector of `params`. This struct stores
+information about that range, as well as whether the sub-vector represents a linked value or
+an unlinked value.
+
+$(TYPEDFIELDS)
+"""
+struct RangeAndLinked
+    # indices that the variable corresponds to in the vectorised parameter
+    range::UnitRange{Int}
+    # whether it's linked
+    is_linked::Bool
+end
+
+"""
+    VectorWithRanges(
+        iden_varname_ranges::NamedTuple,
+        varname_ranges::Dict{VarName,RangeAndLinked},
+        vect::AbstractVector{<:Real},
+    )
+
+A struct that wraps a vector of parameter values, plus information about how random
+variables map to ranges in that vector.
+
+In the simplest case, this could be accomplished only with a single dictionary mapping
+VarNames to ranges and link status. However, for performance reasons, we separate out
+VarNames with identity optics into a NamedTuple (`iden_varname_ranges`). All
+non-identity-optic VarNames are stored in the `varname_ranges` Dict.
+
+It would be nice to improve the NamedTuple and Dict approach. See, e.g.
+https://github.com/TuringLang/DynamicPPL.jl/issues/1116.
+"""
+struct VectorWithRanges{N<:NamedTuple,T<:AbstractVector{<:Real}}
+    # This NamedTuple stores the ranges for identity VarNames
+    iden_varname_ranges::N
+    # This Dict stores the ranges for all other VarNames
+    varname_ranges::Dict{VarName,RangeAndLinked}
+    # The full parameter vector which we index into to get variable values
+    vect::T
+end
+
+function _get_range_and_linked(
+    vr::VectorWithRanges, ::VarName{sym,typeof(identity)}
+) where {sym}
+    return vr.iden_varname_ranges[sym]
+end
+function _get_range_and_linked(vr::VectorWithRanges, vn::VarName)
+    return vr.varname_ranges[vn]
+end
+function init(
+    ::Random.AbstractRNG,
+    vn::VarName,
+    dist::Distribution,
+    p::InitFromParams{<:VectorWithRanges},
+)
+    vr = p.params
+    range_and_linked = _get_range_and_linked(vr, vn)
+    transform = if range_and_linked.is_linked
+        from_linked_vec_transform(dist)
+    else
+        from_vec_transform(dist)
+    end
+    return (@view vr.vect[range_and_linked.range]), transform
+end
+
 """
     AbstractInitStrategy
 
@@ -194,74 +262,6 @@ function init(
     end
 end
 
-"""
-    RangeAndLinked
-
-Suppose we have vectorised parameters `params::AbstractVector{<:Real}`. Each random variable
-in the model will in general correspond to a sub-vector of `params`. This struct stores
-information about that range, as well as whether the sub-vector represents a linked value or
-an unlinked value.
-
-$(TYPEDFIELDS)
-"""
-struct RangeAndLinked
-    # indices that the variable corresponds to in the vectorised parameter
-    range::UnitRange{Int}
-    # whether it's linked
-    is_linked::Bool
-end
-
-"""
-    VectorWithRanges(
-        iden_varname_ranges::NamedTuple,
-        varname_ranges::Dict{VarName,RangeAndLinked},
-        vect::AbstractVector{<:Real},
-    )
-
-A struct that wraps a vector of parameter values, plus information about how random
-variables map to ranges in that vector.
-
-In the simplest case, this could be accomplished only with a single dictionary mapping
-VarNames to ranges and link status. However, for performance reasons, we separate out
-VarNames with identity optics into a NamedTuple (`iden_varname_ranges`). All
-non-identity-optic VarNames are stored in the `varname_ranges` Dict.
-
-It would be nice to improve the NamedTuple and Dict approach. See, e.g.
-https://github.com/TuringLang/DynamicPPL.jl/issues/1116.
-"""
-struct VectorWithRanges{N<:NamedTuple,T<:AbstractVector{<:Real}}
-    # This NamedTuple stores the ranges for identity VarNames
-    iden_varname_ranges::N
-    # This Dict stores the ranges for all other VarNames
-    varname_ranges::Dict{VarName,RangeAndLinked}
-    # The full parameter vector which we index into to get variable values
-    vect::T
-end
-
-function _get_range_and_linked(
-    vr::VectorWithRanges, ::VarName{sym,typeof(identity)}
-) where {sym}
-    return vr.iden_varname_ranges[sym]
-end
-function _get_range_and_linked(vr::VectorWithRanges, vn::VarName)
-    return vr.varname_ranges[vn]
-end
-function init(
-    ::Random.AbstractRNG,
-    vn::VarName,
-    dist::Distribution,
-    p::InitFromParams{<:VectorWithRanges},
-)
-    vr = p.params
-    range_and_linked = _get_range_and_linked(vr, vn)
-    transform = if range_and_linked.is_linked
-        from_linked_vec_transform(dist)
-    else
-        from_vec_transform(dist)
-    end
-    return (@view vr.vect[range_and_linked.range]), transform
-end
-
 """
     InitContext(
             [rng::Random.AbstractRNG=Random.default_rng()],