JuliaDiff / JuliaDiff/ChainRulesCore.jl
variable scope of `@thunk`
Nobody has claimed this yet.
- Dominant language
- Julia
- Stars
- 267
- Forks
- 66
- PR merge metrics
- No merged PRs in 30d
Description
#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
@asyncand@spawnusing$to interpolate values. so@thunkwould 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.
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by reviewing the @thunk macro, _usethunks() path, and issue #683 to understand the intended behavior for second-order AD. Reproduce the caller-scope and Core.Box type-instability issue, then verify the relevant rrule type-stability checks in NeuralAttentionlib.jl while preserving correct thunk behavior in both branches.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- julia
- Domain
- backend, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 32/100