Skip to content

variable scope of @thunk #684

@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
    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions