[Parvm1102]

The findnz pullbacks destructure the incoming tangent and pass the values component straight to sparse/sparsevec, but that component can itself be a thunk — unthunk on the outer tangent doesn't unthunk its components. For example:

using SparseArrays, Zygote
S = sprand(5, 5, 0.4)
Zygote.gradient(A -> sum(findnz(A)[3]), S)
# ERROR: MethodError: no method matching sparse(::Vector{Int64}, ::Vector{Int64}, ::InplaceableThunk{...}, ::Int64, ::Int64)

This unthunks the component in both the matrix and vector pullbacks and adds tests that call them with thunked components. With the fix, the gradient above returns the expected SparseMatrixCSC.

[CarloLucibello]

@oxinabox can we merge this? CI failures appear unrelated

[devmotion]

I'm surprised about this issue in the first place - if unthunk wouldn't ensure that thunks are unthunked, we'd have to add recursive unthunk everywhere. The docs also suggest that thunks inside of a thunk are not desired:

Do not use @thunk if this would be equal or more work than actually evaluating the expression itself. Examples being:

• ...
• The expression being itself a thunk

[devmotion]

Based on my gut feeling, I looked into this issue together with Claude. A Tangent may contain thunks in its fields. Zygote produces exactly that here: Tangent{Any}((ZeroTangent(), ZeroTangent(), thunk)). getindex/getproperty on a Tangent unthunk the field they return, but destructuring (_, _, V̄ = Δ) goes through indexed_iterate, which doesn't.

So rather than unthunking V̄ after destructuring, I pushed a commit that makes the pullbacks index the one component they need (unthunk(Δ)[3] / unthunk(Δ)[2]). That unthunks only that component and leaves the others lazy. The tests now pass a proper Tangent instead of a plain tuple, with @thunk(error()) in the unused components, so they also check that those components are never computed.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 96.57%. Comparing base (85be40a) to head (1133868).

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #843   +/-   ##
=======================================
  Coverage   96.57%   96.57%           
=======================================
  Files          28       28           
  Lines        3618     3618           
=======================================
  Hits         3494     3494           
  Misses        124      124           

☔ 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.

[devmotion]

Thank you!