Modify variable in closure and type inference

General Usage
1

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
2

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
3

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
4

Wow that indeed works. I thought this would run into performance of captured variables in closures · Issue #15276 · JuliaLang/julia · GitHub, but it does not. Now I am confused. Can you comment on why this avoids 15276?

5

Any reason to not use mapreduce/mapfoldl for this?

6

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
7

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 https://github.com/JuliaLang/julia/issues/15276#issuecomment-297596373

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

8

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
9

This seems to run into performance of captured variables in closures · Issue #15276 · JuliaLang/julia · GitHub 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
10

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
11

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?

12

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

1 Like