JuliaDiff / JuliaDiff/ChainRules.jl

Pullback for sparse-array vector product very inefficient

Open
#803 4 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Julia
Stars
475
Forks
98
PR merge metrics
No merged PRs in 30d

Description

Our current rrule for sparse matrix vector products is very inefficient, and causes out-of-memory with large sparse CPU or GPU arrays. Our current rrule(*, sparse(A), x) is implemented like this

function rrule(
    ::typeof(*),
    A::AbstractVecOrMat{<:CommutativeMulNumber},
    B::AbstractVecOrMat{<:CommutativeMulNumber},
)
    project_A = ProjectTo(A)
    ...
        dA = @thunk(project_A(Ȳ * B'))
    ...
end

So we first compute a non-sparse Ȳ * B' (may easily exceed memory if A was very large but very sparse) and then project back to a sparse tangent.

The best way to fix this (at least if Ȳ' and 'B' are vectors) might be adding a specific "vector-outer-product" array type for read-only vector * adjoint-vector products (might be useful in general) that computes getindex on the fly. Or maybe we already have that somewhere?

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start at the rrule for sparse matrix–vector products described in the issue and trace how the dA thunk computes Ȳ * B'. Investigate whether an existing lazy vector-outer-product representation can avoid materializing a dense intermediate; done means sparse CPU and GPU cases no longer risk dense-memory allocation while preserving the sparse tangent result.

Written by the indexing model from the issue text.

Assessment

Tech stack
julia
Domain
performance
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.