Help eliminate a captured variable

General Usage Performance
1

I am trying to eliminate a captured variable, as reported by JET.jl, in this piece of code, copied here for convenience:

function transform_with(flag::LogJacFlag, transformation::StaticArrayTransformation{D,S},
                        x::AbstractVector{T}, index::Int) where {D,S,T}
    (; inner_transformation) = transformation
    # NOTE this is a fix for #112, enforcing types taken from the transformation of the
    # first element.
    y1, β„“1, index1 = transform_with(flag, inner_transformation, x, index)
    D == 1 && return SArray{S}(y1), β„“1, index1
    L = typeof(β„“1)
    let β„“::L = β„“1, index::Int = index1
        function _f(_)
            y, β„“Ξ”, indexβ€² = transform_with(flag, inner_transformation, x, index)
            index = indexβ€²
            β„“ = β„“ + β„“Ξ”
            y
        end
        yrest = SVector{D-1}(_f(i) for i in 2:D)
        SArray{S}(pushfirst(yrest, y1)), β„“, index
    end
end

The reproducer is below, on Julia 1.12.5, JET.jl reports

julia> using TransformVariables, JET, StaticArrays

julia> t = as(SVector{3}, asβ„β‚Š)
TransformVariables.StaticArrayTransformation{3, Tuple{3}, TVExp}(asβ„β‚Š) (dimension 3)

julia> @report_opt transform(t, ones(3))
═════ 2 possible errors found ═════
β”Œ transform(t::TransformVariables.StaticArrayTransformation{3, Tuple{3}, TVExp}, x::Vector{Float64}) @ TransformVariables /home/tamas/code/julia/TransformVariables/src/generic.jl:315
β”‚β”Œ transform_with(flag::TransformVariables.NoLogJac, transformation::TransformVariables.StaticArrayTransformation{…}, x::Vector{…}, index::Int64) @ TransformVariables /home/tamas/code/julia/TransformVariables/src/aggregation.jl:210
β”‚β”‚ captured variable `β„“` detected
│└────────────────────
β”‚β”Œ transform_with(flag::TransformVariables.NoLogJac, transformation::TransformVariables.StaticArrayTransformation{…}, x::Vector{…}, index::Int64) @ TransformVariables /home/tamas/code/julia/TransformVariables/src/aggregation.jl:210
β”‚β”‚ captured variable `index` detected
│└────────────────────

I thought that the let would take care of this, what am I missing? What’s the recommended fix?

2

The let approach doesn’t work if you assign the variables within closures. I’d recommend using mutable container to avoid direct assignments in such cases:
https://aviatesk.github.io/JETLS.jl/release/diagnostic/#Workaround-d9291b076db1c42c

5 Likes