Avoid capturing input array x in the chunk pullback - #252
Merged
Conversation
The reverse rule for `chunk` captured the whole input array `x` in its pullback closure, keeping it alive for the entire backward pass even though `chunk`'s gradient only needs `x`'s axes, array type and element type (never its values). Capture `ProjectTo(x)` and an empty `similar(x, 0)` instead: together they supply the axes, structure and array/element type needed to allocate and project the gradient, without holding any of `x`'s data. On a 1000x500 input the pullback's captures shrink from ~4 MB to 80 B. `∇chunk` and its second-order rrule are updated to the new signature, as is the direct `∇chunk` test (the `ProjectTo`/template args are held fixed since they are non-differentiable). Also relax the `CUDA` test compat to `5, 6` so the test environment resolves on current Julia. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Resolves the
# TODO avoid capturing x in the pullbackatsrc/utils.jl:219(item 1 of #227).Problem
rrule(::typeof(chunk), x)captured the whole input arrayxin its pullback closure, keeping it alive for the entire backward pass — even thoughchunk's gradient never usesx's values, only its axes, array type and element type.Change
Capture
ProjectTo(x)and an emptysimilar(x, 0)instead ofx:ProjectTo(x)carriesx's axes and restores its structure.similar(x, 0)carriesx's array/element type (so the gradient allocates with the right backend, e.g. GPU) while holding no data.∇chunkbuildsdxfrom these and returnsproject(dx). The second-orderrrule(::typeof(∇chunk))and the direct∇chunktest are updated to the new signature.On a 1000×500 input, the pullback's captured state shrinks from ~4 MB to 80 bytes; the all-zero branch still returns a materialized zero array, so type stability (checked by the
size collectiontest) is preserved.Also relaxes the
CUDAtest compat to5, 6so the test environment resolves on current Julia.Testing
The full
chunktestset passes (37/37) in the repo's test environment.CUDA.functional()wasfalselocally, so the GPU branch was not exercised.🤖 Generated with Claude Code