Skip to content

Disable thunks for 2nd order AD#683

Merged
mcabbott merged 6 commits into
JuliaDiff:mainfrom
pxl-th:pxl-th/thunks
Jan 3, 2025
Merged

Disable thunks for 2nd order AD#683
mcabbott merged 6 commits into
JuliaDiff:mainfrom
pxl-th:pxl-th/thunks

Conversation

@pxl-th

@pxl-th pxl-th commented Jan 2, 2025

Copy link
Copy Markdown
Contributor

Needed for lazy Zygote:
FluxML/Zygote.jl#966

Also bumped patch version so that we can tag a release afterwards.

@pxl-th pxl-th changed the title Allow disabling thunking Disable thunks for 2nd order AD Jan 2, 2025
Comment thread src/tangent_types/thunks.jl Outdated
Comment on lines +148 to +152
return quote
$(esc(_usethunks))() ?
Thunk($(esc(func))) :
$(esc(func))()
end

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Comments on this approach over :(_usethunks() ? Thunk($(esc(func))) : $(esc(body))), from #568? My hope there was that including the body directly, instead of making and calling a function, might be slightly simpler for e.g. Zygote to understand.

@pxl-th pxl-th requested a review from mcabbott January 2, 2025 23:09

@mcabbott mcabbott left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's do this.

CI failures look unrelated

Comment on lines 145 to 147
# Basically `:(Thunk(() -> $(esc(body))))` but use the location where it is defined.
# so we get useful stack traces if it errors.
func = Expr(:->, Expr(:tuple), Expr(:block, __source__, body))

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This could go further and generate a name for the thunk, based on source location, instead of anon. () ->. I think the name will be seen sometimes where the code line info is not?

But not this PR perhaps...

@mcabbott mcabbott merged commit 2aad117 into JuliaDiff:main Jan 3, 2025
@pxl-th pxl-th deleted the pxl-th/thunks branch January 3, 2025 06:54
@oxinabox

oxinabox commented Jan 3, 2025

Copy link
Copy Markdown
Member

This was not a patch change.
This was not a bug.

This is a new feature.
This should have been a new minor version.

I recommended being careful about this and biasing towards doing minor releases because it makes it possible to do backports in between releases
(mostly addressed at @mcabbott).

Not a huge deal as we are unlikely to need to put in patches.
But it is always not a problem until it is.

@pxl-th

pxl-th commented Jan 3, 2025

Copy link
Copy Markdown
Contributor Author

Fair point. Initially this PR added a hook that users could register that would allow to disable thunking and so the original behavior was not changed, hence the patch version.

But then I changed the PR to disable thunking for 2nd order AD unconditionally and forgot to change the version bump from patch to minor.

@oxinabox

oxinabox commented Jan 6, 2025

Copy link
Copy Markdown
Member

That would still have been a new feature and would have been a minor.
Patch is exclusively for fixing bugs.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants