diff --git a/src/Strat/params.jl b/src/Strat/params.jl index b45321cb5c5b1dafe92405cf0825fbdb19cad46e..2e9db5a7d63186fb3bd82962f4b2a48d4589834b 100644 --- a/src/Strat/params.jl +++ b/src/Strat/params.jl @@ -31,9 +31,10 @@ ComponentArrays.ComponentArray(rv::ParameterVector) = getfield(rv, :vals) function parameters(model::Tile, transforms::Pair{Symbol,<:Pair{Symbol,<:ParamTransform}}...) getparam(x) = x - getparam(x::Union{<:AbstractVector,<:Tuple}) = length(x) == 1 ? getparam(x[1]) : collect(getparam.(x)) + getparam(x::Union{<:AbstractVector,<:Tuple}) = length(x) == 1 ? getparam(x[1]) : Tuple(getparam.(x)) paramval(x) = ustrip(x) paramval(x::Param) = ustrip(x.val) + paramval(x::Tuple) = collect(x) m = Model(model) nestedparams = mapflat(getparam, groupparams(m, :layer, :fieldname); maptype=NamedTuple) mappedparams = nestedparams @@ -44,7 +45,7 @@ function parameters(model::Tile, transforms::Pair{Symbol,<:Pair{Symbol,<:ParamTr @set! mappedparams[layer][var] = mapflat(getparam, groupparams(m_transform, :transform, :fieldname); maptype=NamedTuple) push!(mappings, ParamMapping(transform, var, layer)) end - mappedarr = ComponentArray(mapflat(paramval, mappedparams)) + mappedarr = ComponentArray(mapflat(paramval, mappedparams; maptype=NamedTuple)) return ParameterVector(mappedarr, mappedparams, mappings...) end @inline @generated function updateparams!(v::AbstractVector, model::Tile, u, du, t)