Skip to content

Unthunk cotangent components in findnz pullbacks - #843

Open
Parvm1102 wants to merge 1 commit into
JuliaDiff:mainfrom
Parvm1102:fix-findnz-thunked-cotangent
Open

Unthunk cotangent components in findnz pullbacks#843
Parvm1102 wants to merge 1 commit into
JuliaDiff:mainfrom
Parvm1102:fix-findnz-thunked-cotangent

Conversation

@Parvm1102

Copy link
Copy Markdown

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.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
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