Is there a way to modify captured variables in a closure?

Suppose I have some closure like,

f = let x = 2
    () -> x^2
end

Is there some way to create a new closure from this where I modify the type and value of x? Something like,

g = deepcopy(f)
g.x = "2"
g() # would give "22"

The above does not work since you can’t set g.x, but gives you the idea of what I’m looking for. The reason I want this is basically because I have some closures with captured variables that are CuArrays that I’m trying to serialize, which later get deserialized on a machine with no GPU, so I want a way to pre-convert them to Arrays, but keep the original closure function as-is.

1 Like

I don’t think this is easily possible. But do you need to do this with a closure or would a functor also work?
E.g.

julia> struct A
           x::Float64
       end

julia> struct B
           x::Int
       end

julia> (o::A)() = o.x^2    # or (o::Union{A,B})() = o.x^2 in julia 1.3.0

julia> (o::B)() = o.x^2

julia> Base.convert(::Type{B}, o::A) = B(o.x)

julia> a = A(12)
A(12.0)

julia> a()
144.0

julia> convert(B, a)()
144

Edit: sorry this is not what you asked…

Like so:

julia> let x=0
       global inc() = (x+=1; nothing)
       global f(a) = x+a
       end
f (generic function with 1 method)

julia> f(1)
1

julia> inc()

julia> f(1)
2

According to the devdocs, your code is almost identical to the code that defining the closure in the first place “roughly lowers to”, so I’m just hopeful that there’s a way to use that and avoid the boilerplate that results from writing it out by hand.

That’s interesting, I hadn’t realized you could do that. It kind of gives me a solution, although I think unfortunately I’d have to write “setter” functions for all my captured variables which I think which mostly defeats the purpose.

f = let x = 2
    global setx(newx) = (x=newx)
    () -> x^2
end

f() # gives 4

setx("2")

f() # gives "22"

You could make a factory:

julia> make_closure(x) = () -> x^2                                                                                                                                       
make_closure (generic function with 1 method)                                                                                                                            
                                                                                                                                                                        
julia> f1 = make_closure(2)                                                                                                                                              
#9 (generic function with 1 method)                                                 

julia> f1()    
4       
                                          
julia> f2 = make_closure("x")
#9 (generic function with 1 method)
                                          
julia> f2()
"xx"               

I think that (or the solution by @jbrea) would be better than messing with the closure’s internals.

1 Like

Yes, but you have control over the naming of the structures, which is handy when you want to define the conversion. Does it add so much boilerplate code to write it like this? If you can use a parametric type it becomes only

julia> struct A{T}
           x::T
       end

julia> (a::A)() = a.x^2

julia> Base.convert(::Type{A{T}}, a::A) where T = A{T}(a.x)

julia> a = A(12.)
A{Float64}(12.0)

julia> a()
144.0

julia> convert(A{Int}, a)()
144

julia> A("x")()
"xx"

Thanks for the suggestions. I think I’ve got a solution that just uses the closure itself. Given that making the closure supposedly lowers to that struct, I was kind of confused why e.g. typeof(f)(3) wouldn’t let me construct a new instance of the closure where x was 3. I still don’t know why it doesn’t, but you can construct it by creating a “new” expression by hand and eval’ing it:

f = let x = 2
    () -> x^2
end

g = eval(Expr(:new, typeof(f), 3))

g() # gives 9

g = eval(Expr(:new, typeof(f).name.wrapper{String}, "2"))

g() # gives "22"

And you can stick this in a @generated function that can then programatically scan through the closed over variables and convert the ones that are CuArrays.

Although a followup question I now have is why does eval(Expr(:new, typeof(f), 3)) work but not typeof(f)(3)? It would be nice if it did so I could avoid @generated world-age issues.

Sounds fancy :smile:. But what does really speak against the simple suggestion I posted above?

Also eval will evaluate in global scope.

Interesting question.

I believe this is because the struct that the closure is lowered to doesn’t have any default constructors. They are not normally required; lowering just uses a new expression to create the closure in situ.

1 Like

Makes sense, I hadn’t appreciated the default constructor doesn’t always come for free.

Yes, but using a @generated function gets rid of that.

Nothing at all, its just that my real problem I have alot of these closures with many closed over variables all with different names, etc… so I was just looking for something more automatic.

FWIW, here’s more or less what the final solution looks like, here modifying an arbitrary number of closed over variables. I’m pretty impressed with Julia that its this simple and also completely type stable:

@generated function set_captured_vars(f::F, vals...) where {F<:Function}
    Expr(:splatnew, F.name.wrapper{vals...}, :vals)
end

f = let x=2, y=3
    () -> (x^2, y)
end

f() # gives (4,3)

g = set_captured_vars(f, "2", 3.)

g() # gives ("22", 3.00)
2 Likes

@oxinabox I believe you were looking for something like this recently.

1 Like

Yeah,
I was wanting a constructor for closures.
Cos they are basically functors (* kinda.)
If I had an instance I want to make a new instance with different closed variables.

I know it is possible,
but all I have is notes saying “Find the ccall that BSON.jl / Serialize uses to instanstiate them”.
Jeff says that is legit,
and the only reasom they don’t is to save time on creating the new methods.
IIRC it is ccall(:jl_newstructv or something like that.

But maybe it would be easier to use @marius311’s @generated function

Late to the party, but this is how it seems to work with jl_new_structtt:

f = let x=2, y=3
    () -> (x^2, y)
end

f() # gives (4,3)

flds = ["2", 3.]
g = ccall(:jl_new_structt, Any, (Any, Any),
          typeof(f).name.wrapper{typeof.(flds)...}, Tuple(flds))

g() # gives ("22", 3.00)

I found a case where this doesn’t work though, because the closure struct has additional typevars (from static parameters):

function main(A::Array{T}=zeros(1), b::Number=0) where {T}
    function f(A)
        convert(T, 0)
        A[1] = b
        return
    end

    @show fieldnames(typeof(f)) # (:b,)
    @show typeof(f).types       # svec(Int64)
    @show typeof(f).parameters  # svec(Float64, Int64)

    # replace captured variable :b
    flds = [2]
    f = ccall(:jl_new_structt, Any, (Any, Any),
              typeof(f).name.wrapper{typeof.(flds)...}, Tuple(flds))

    f(A)

    A
end

Apparently those always come first, so it’s safe to skip those. That’s what I’m doing here: https://github.com/JuliaGPU/CUDA.jl/pull/625/commits/5ed800695061b3a95b644f02fee50511ce860ce5

3 Likes

I’ve added this functionality to Adapt.jl, so for the purpose of OP it should be as easy as calling adapt:

julia> using CUDA, Adapt

julia> function foo(A)
       bar() = A
       bar
       end
foo (generic function with 1 method)

julia> f = foo(rand(2,2))
(::var"#bar#1"{Matrix{Float64}}) (generic function with 1 method)

julia> f()
2×2 Matrix{Float64}:
 0.167181  0.62589
 0.174081  0.464358

julia> g = Adapt.adapt(CuArray, f)
(::var"#bar#1"{CuArray{Float64, 2}}) (generic function with 1 method)

julia> g()
2×2 CuArray{Float64, 2}:
 0.167181  0.62589
 0.174081  0.464358
1 Like