Skip to content

Make the getindex add! agree with its value - #872

Open
CarloLucibello wants to merge 1 commit into
mainfrom
cl/getindex-add
Open

CarloLucibello wants to merge 1 commit into
mainfrom
cl/getindex-add

Conversation

@CarloLucibello

Copy link
Copy Markdown
Member

Part 1 of #863. The add! of the InplaceableThunk from the getindex/view rules added the raw dy, skipping the ProjectTo(x) that val applies. So t.add!(dx) could differ from dx + unthunk(t), or throw (e.g. a complex dy for a real x), and an AD couldn't safely call it.

  • When ProjectTo(x) acts elementwise (arrays of numbers, including GPU arrays), add! projects dy and adds it into dx[inds...], staying O(length(dy)). Arrays of arrays use only the projectors of the indexed elements. Other arrays (Diagonal, Symmetric, sparse, …) fall back to add!!(dx, val).
  • Scalar ∇getindex projects dy before building the OneElement, so a Float64 cotangent for a Float32 array no longer materializes a dense array.

Benchmarks with Zygote 0.7.16 on Julia 1.13:

  • x[i] loop, x::Vector{Float32} of length 10⁴, summed in Float64: 26.5 ms / 391 MiB → 1.80 ms / 8.3 MiB.
  • X[:, j] loop over a 100×1000 matrix: unchanged at 36 ms / 767 MiB, since Zygote unthunks. With Zygote calling add! when the inferred eltype of val fits its buffer, which this PR makes safe: 0.41 ms / 5.1 MiB.
  • add! itself costs the same as before, on CPU and on GPU.

🧔‍♂️🤖 Assisted by AI

The `add!` of the `InplaceableThunk` returned by the `getindex` and `view`
rules added the raw cotangent, skipping the `ProjectTo(x)` that `val` applies,
so it could differ from `unthunk` or throw. It now projects `dy` when
`ProjectTo(x)` acts elementwise, staying O(length(dy)), and adds `val`
otherwise. Scalar `∇getindex` also projects `dy` before building the
`OneElement`, which avoids materializing it for mixed precision.

Part 1 of #863.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Comment on lines +2 to +3
ChainRulesCore, ChainRulesTestUtils, FiniteDifferences, GPUArraysCore, JLArrays,
LinearAlgebra, Test

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
ChainRulesCore, ChainRulesTestUtils, FiniteDifferences, GPUArraysCore, JLArrays,
LinearAlgebra, Test
ChainRulesCore,
ChainRulesTestUtils,
FiniteDifferences,
GPUArraysCore,
JLArrays,
LinearAlgebra,
Test

@codecov

codecov Bot commented Nov 11, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 88.23529% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 96.61%. Comparing base (d9ef0f7) to head (6ac471b).

Files with missing lines Patch % Lines
src/rulesets/Base/indexing.jl 88.23% 2 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main     #872      +/-   ##
==========================================
- Coverage   96.65%   96.61%   -0.05%     
==========================================
  Files          28       28              
  Lines        3615     3630      +15     
==========================================
+ Hits         3494     3507      +13     
- Misses        121      123       +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.

This branch has not been deployed

No deployments
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.

1 participant