Skip to content

variable scope of @thunk #684

Open
@chengchingwen

Description

@chengchingwen

#683 introduces a separate path when creating the thunk (_usethunks() ? Thunk($(esc(func))) : $(esc(body))). However, this would introduce new variables into the caller's scope and create Core.Box thus causing type instability, which breaks some rrule type stability checks in NeuralAttentionlib.jl.

Currently, I have two ways to solve this:

  1. replace $(esc(body)) with $(esc(func))() which turns the else path into a function call, thus no variable introduction.
  2. mimic how @async and @spawn using $ to interpolate values. so @thunk would be defined as:
macro thunk(body)
    letargs = Base._lift_one_interp!(body)
    func = Base.replace_linenums!(:(()->($(esc(body)))), __source__)
    return quote
        if _usethunks()
            let $(letargs...)
                Thunk($func)
            end
        else
            let $(letargs...)
                $(esc(body))
            end
        end
    end
end

However, I don't understand #683 enough to see what would be better for 2nd order AD. Maybe @pxl-th or @mcabbott would have some thoughts on this.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions