Modify variable in closure and type inference

I would like to have a closure modify a captured variable. Here is a simple example:

function fancy_iteration(f::F, itr) where {F}
    for pt in itr
        f(pt)
    end
end

function mysum(f, itr)
    sum = 0.
    fancy_iteration(itr) do pt
        sum += f(pt)
    end
    sum
end

@code_warntype mysum(sin, ones(3))

Currently inference fails on this. Are the recommended workarounds?

1 Like

You can use a Ref, like this:

julia> function mysum(f, itr)
         sum = Ref(0.0)
         fancy_iteration(itr) do pt
           sum[] += f(pt)
         end
         sum[]
       end
mysum (generic function with 1 method)
4 Likes

Why don’t you just return the new value from the function and modify the variable with the new value outside the function call?

2 Likes

Wow that indeed works. I thought this would run into https://github.com/JuliaLang/julia/issues/15276, but it does not. Now I am confused. Can you comment on why this avoids 15276?

Any reason to not use mapreduce/mapfoldl for this?

You are no longer assigning to sum (you are just updating the content of the container) which means that there is no need to box the variable.

2 Likes

Ah so 15276 only happens when I assign variables? I remember having problems with read only, too. But maybe that was fixed?
Edit: Here is a read only example taken from performance of captured variables in closures · Issue #15276 · JuliaLang/julia · GitHub

julia> @noinline function call_func(f::Function)
           println(f())
       end
call_func (generic function with 1 method)

julia> # Simple version.  `i` is a Box here.
       function log_loop(N, unlikely)
           s = 0
           for i=1:N
               s += i
               if unlikely
                   call_func(()->"asdf $i")
               end
           end
       end
log_loop (generic function with 1 method)

julia> @code_warntype log_loop(1, true) # fails to infer `i`

How is this different from

Yeah sure in this example one can even use the sum function. But the goal of the question is to be as close to the style of the example, while avoiding inference problems.

1 Like

This seems to run into https://github.com/JuliaLang/julia/issues/15276 in more complicated situations. The following fails for me:

function fancy_iteration(f, itr)
    for pt in itr
        f(pt)
    end
end

function mysum2(f, itr)
    for i in 1:1
        sum = Ref(0.0)
        if true
            fancy_iteration(itr) do pt
                old = sum[]
                new = old + f(pt)
                sum[] = new
            end
        end
        return sum[]
    end
    error()
end

@code_warntype mysum2(sin, 1:3) #fails

The traditional let-block workaround works just fine, though:

julia> function mysum2(f, itr)
           for i in 1:1
               sum = Ref(0.0)
               if true
                   let sum = sum
                       fancy_iteration(itr) do pt
                           old = sum[]
                           new = old + f(pt)
                           sum[] = new
                       end
                   end    
               end
               return sum[]
           end
           error()
       end

You can think of the let block as a really obvious promise to the compiler that nothing inside the closure could possibly modify what sum is bound to in the outer scope.

3 Likes

Thanks I also just tried that! Would you say that inference does well if I use let blocks + Ref for updating captured variables or are there other traps?

Yes, that should work quite well as far as I know.

1 Like