Preserve Tracker gradients in GPUArraysCore restructure - #505
Conversation
`restructure` on a GPUArraysCore target adapted the TrackedArray source through `Tracker.adapt_structure`, which calls `param()` and starts a new tape leaf. The forward value stayed tracked but the source gradient was silently zero. Reshape the TrackedArray instead, matching the existing Array/TrackedArray methods. Fixes JuliaArrays#504 Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Grok <noreply@x.ai> Agent-Harness: Grok CLI 1.0.13 Agent-Model: grok-4.6 Agent-Session: 01a05298-11d9-7a71-89a3-e977fe6a9a3e
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #505 +/- ##
==========================================
+ Coverage 60.93% 61.35% +0.42%
==========================================
Files 15 16 +1
Lines 622 634 +12
==========================================
+ Hits 379 389 +10
- Misses 243 245 +2 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
The It is A separate investigation is running in parallel to bisect/fix or file that master failure. |
A reshape-only path preserved the tape but left a CPU TrackedArray on the host when the source was not already a GPUArraysCore array. Adapt through a tracked primitive when the data is not already the target type; skip Adapt when it is, since Tracker.adapt_structure would start a new tape leaf. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Grok <noreply@x.ai> Agent-Harness: Grok CLI 1.0.13 Agent-Model: grok-4.6 Agent-Session: 01a05298-11d9-7a71-89a3-e977fe6a9a3e
|
Follow-up: a reshape-only path kept the tape but left a CPU CPU source onto a JLArray target now returns JLArray data with an all-ones gradient; GPU source still reshapes in place. |
Fixes #504
What changed and why
ArrayInterface.restructureon aGPUArraysCore.AbstractGPUArraytarget currently doesAdapt.adapt(parameterless_type(x), y)thenreshape. Tracker'sadapt_structureisparam(adapt(T, data(xs))), which starts a new tape leaf. The forward result stays numerically correct and still looks tracked, but the source gradient is silently zero.This adds a
GPUArraysCore+Trackerpackage extension that reshapes theTrackedArrayinstead, matching the existing Tracker methods forArrayandTracker.TrackedArraytargets from #498. Untracked sources still go through Adapt so CPU→GPU conversion is unchanged.Failing before / passing after
The new JLArray test (no GPU hardware) on unfixed
restructure:With the extension:
Verification
JuliaFormatter SciML style and
typoswere clean on the diff.git diff --checkwas clean.Not verified locally
AbstractGPUArray; not in scope.Reviewer notes
TrackedArrayrestructured onto a GPU template stays a CPU tracked array. That matchesrestructure(::Array, ::TrackedArray)and is what preserves the tape. Untrackedystill Adapts onto the GPU type.JLArraysis MIT (from GPUArrays.jl).7.30.0→7.30.1.Links
🤖 Generated with Grok CLI 1.0.13 (model: grok-4.6)
Session: 01a05298-11d9-7a71-89a3-e977fe6a9a3e (local Grok CLI session; no shareable URL)