#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:
- replace
$(esc(body)) with $(esc(func))() which turns the else path into a function call, thus no variable introduction.
- 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.
#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 createCore.Boxthus causing type instability, which breaks some rrule type stability checks in NeuralAttentionlib.jl.Currently, I have two ways to solve this:
$(esc(body))with$(esc(func))()which turns the else path into a function call, thus no variable introduction.@asyncand@spawnusing$to interpolate values. so@thunkwould be defined as: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.