Skip to content

Commit

Permalink
Don't create nested thunks when accumulating
Browse files Browse the repository at this point in the history
Otherwise, it's too easy to create massive types that freeze compilation and blow the stack.
  • Loading branch information
ToucheSir authored Jan 24, 2025
1 parent 572eb2a commit 44b05ab
Showing 1 changed file with 3 additions and 3 deletions.
6 changes: 3 additions & 3 deletions src/lib/lib.jl
Original file line number Diff line number Diff line change
Expand Up @@ -43,9 +43,9 @@ accum(x::ChainRulesCore.Tangent, y::NamedTuple) = accum(wrap_chainrules_output(x
accum(x::Nothing, y::AbstractThunk) = y
accum(x::AbstractThunk, y::Nothing) = x

accum(x, y::AbstractThunk) = @thunk(accum(x, unthunk(y)))
accum(x::AbstractThunk, y) = @thunk(accum(unthunk(x), y))
accum(x::AbstractThunk, y::AbstractThunk) = @thunk(accum(unthunk(x), unthunk(y)))
accum(x, y::AbstractThunk) = accum(x, unthunk(y))
accum(x::AbstractThunk, y) = accum(unthunk(x), y))
accum(x::AbstractThunk, y::AbstractThunk) = accum(unthunk(x), unthunk(y))

# Core functions
@_adjoint_keepthunks deepcopy(x) = deepcopy(x), ȳ -> (ȳ,)
Expand Down

0 comments on commit 44b05ab

Please sign in to comment.