Skip to content

Preserve Tracker gradients in GPUArraysCore restructure - #505

Merged
ChrisRackauckas merged 2 commits into
JuliaArrays:masterfrom
ChrisRackauckas-Claude:restructure-tracker-gpu
Aug 30, 2026
Merged

Preserve Tracker gradients in GPUArraysCore restructure#505
ChrisRackauckas merged 2 commits into
JuliaArrays:masterfrom
ChrisRackauckas-Claude:restructure-tracker-gpu

Conversation

@ChrisRackauckas-Claude

Copy link
Copy Markdown
Contributor

Ignore this PR until it has been reviewed by @ChrisRackauckas.

Fixes #504

What changed and why

ArrayInterface.restructure on a GPUArraysCore.AbstractGPUArray target currently does Adapt.adapt(parameterless_type(x), y) then reshape. Tracker's adapt_structure is param(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+Tracker package extension that reshapes the TrackedArray instead, matching the existing Tracker methods for Array and Tracker.TrackedArray targets 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:

restructure GPUArraysCore + Tracker: Test Failed at test/ad.jl:69
  Expression: Array(Tracker.data(dx)) == ones(Float32, 6)
   Evaluated: Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0] == Float32[1.0, 1.0, 1.0, 1.0, 1.0, 1.0]

Test Summary:                       | Pass  Fail  Total  Time
restructure GPUArraysCore + Tracker |    4     1      5  7.8s

With the extension:

Test Summary:                       | Pass  Total  Time
restructure GPUArraysCore + Tracker |    5      5  4.7s

Verification

$ GROUP=Core timeout 3600 julia --project=. -e 'using Pkg; Pkg.test()'
Test Summary:       | Pass  Total  Time
BandedMatrices      |   21     21  3.7s
BlockBandedMatrices |    8      8  5.9s
Core                |  221    221 27.8s
AD Integration      |   33     33 17.2s
StaticArrays        |   39     39  8.5s
ChainRules          |   20     20 22.9s
FillArrays          |    7      7  0.0s
Testing ArrayInterface tests passed

JuliaFormatter SciML style and typos were clean on the diff. git diff --check was clean.

Not verified locally

  • GPU CI group (CUDA/Metal). This host has no NVIDIA device; the JLArray path is the device-independent stand-in from the issue.
  • Downstream SciMLSensitivity TrackerAdjoint GPU job. The workaround in Preserve Tracker array parameter gradients SciML/SciMLSensitivity.jl#1636 should become unnecessary once this is registered, but that was not re-run here.
  • Docs build: no docstring or public API change.
  • ReverseDiff+GPU. ReverseDiff rarely wraps AbstractGPUArray; not in scope.

Reviewer notes

  • The new method reshapes without Adapt, so a CPU TrackedArray restructured onto a GPU template stays a CPU tracked array. That matches restructure(::Array, ::TrackedArray) and is what preserves the tape. Untracked y still Adapts onto the GPU type.
  • Test extra JLArrays is MIT (from GPUArrays.jl).
  • Patch bump 7.30.07.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)

`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

codecov Bot commented Aug 30, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 66.66667% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 61.35%. Comparing base (995c24d) to head (a5bc725).
⚠️ Report is 3 commits behind head on master.

Files with missing lines Patch % Lines
ext/ArrayInterfaceGPUArraysCoreTrackerExt.jl 66.66% 4 Missing ⚠️
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.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@ChrisRackauckas-Claude

Copy link
Copy Markdown
Contributor Author

The test (alldeps, 1.10, Core) Downgrade failure is pre-existing on master, not introduced here.

It is UndefVarError: OneElement not defined while precompiling ArrayInterfaceChainRulesExt (using ChainRules: OneElement) under lowest-compat deps. The same Downgrade workflow has been red on every recent master push, including https://github.com/JuliaArrays/ArrayInterface.jl/actions/runs/32094022221 (merge of #503).

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
@ChrisRackauckas-Claude

Copy link
Copy Markdown
Contributor Author

Follow-up: a reshape-only path kept the tape but left a CPU TrackedArray on the host. restructure now Adapts onto the GPU target type through a tracked primitive when the data is not already that type, and still skips Adapt.adapt on the TrackedArray itself (that goes through param() and zeros the source gradient).

CPU source onto a JLArray target now returns JLArray data with an all-ones gradient; GPU source still reshapes in place.

@ChrisRackauckas
ChrisRackauckas marked this pull request as ready for review August 30, 2026 13:15
@ChrisRackauckas
ChrisRackauckas merged commit 97b9fe8 into JuliaArrays:master Aug 30, 2026
19 of 21 checks passed
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.

restructure drops Tracker gradients for GPUArraysCore targets

2 participants