rrule for broadcasted cast of sparse array

[CarloLucibello]

This also relaxes a bit the signature for sparse constructors.

Tests fail, but maybe is just a matter of passing an appropriate output tanget to test_rrule?

Close FluxML/Zygote.jl#810

[CarloLucibello]

I think that the implementations of both features are correct but

• I don't know how to test the broadcasted rule. Maybe just test with Zygote via test_rrule(..., rrule_f = ZYGOTE_SOMETHING)?
• I cannot manage to get the tests for findnz to work, I get a very long FiniteDifference error

Sorry for adding 2 loosely related features in the same PR, I've been a bit lazy

[mcabbott]

I think JuliaDiff/ChainRulesTestUtils.jl#234 has some hints for testing broadcasted rules.

Although the rules aren't doing any mathematics, so perhaps there's not much value from finite differencing. CRTU does also try to test various types of tangent etc, but maybe this isn't working.

[CarloLucibello]

Pfff, why Julia 1.6 is failing now on an old test

[CarloLucibello]

Should be ready now

[mzgubic]

Thanks! The findnz part LGTM, but I didn't check the broadcasting

[mzgubic]

Actually, for testing, we need the compat for FiniteDifferences to be set to 0.12.23

[mcabbott]

I think this is the first broadcast rule in here. I'm not entirely sure how I feel about this.

T.(x) is eager, would there be downsides to leaving it lazy, like broadcasted(T, x)? That may mean more things fuse, or it may mean that subsequent Zygote broadcasting rules fail because they want an array?

Similarly proj(Δ) won't accept something like a broadcasted, if we explore fused gradients that way.

These are the sort of questions that make me unsure we want to commit to one design here. I guess having one or two won't hurt too much, we can fix them here later if needed.

Zygote's present rule fails, I believe it's trying to make a sparsearray of tuples?

julia> gradient(x -> sum(Float32.(x)), sprand(10, 0.5))
ERROR: MethodError: no method matching zero(::Type{Any})
Closest candidates are:
  zero(::Type{Union{Missing, T}}) where T at ~/.julia/dev/julia/usr/share/julia/base/missing.jl:105
  zero(::Union{Type{P}, P}) where P<:Dates.Period at ~/.julia/dev/julia/usr/share/julia/stdlib/v1.8/Dates/src/periods.jl:53
  zero(::CartesianIndex{N}) where N at ~/.julia/dev/julia/usr/share/julia/base/multidimensional.jl:106
  ...
Stacktrace:
  [1] zero(#unused#::Type{Any})
    @ Base ./missing.jl:106
  [2] _zeros_eltypes
    @ ~/.julia/dev/julia/usr/share/julia/stdlib/v1.8/SparseArrays/src/higherorderfns.jl:208 [inlined]
  [3] _noshapecheck_map(::typeof(first), ::SparseVector{Any, Int64})
    @ SparseArrays.HigherOrderFns ~/.julia/dev/julia/usr/share/julia/stdlib/v1.8/SparseArrays/src/higherorderfns.jl:164
  [4] map(f::typeof(first), A::SparseVector{Any, Int64})
    @ SparseArrays.HigherOrderFns ~/.julia/dev/julia/usr/share/julia/stdlib/v1.8/SparseArrays/src/higherorderfns.jl:147
  [5] adjoint
    @ ~/.julia/packages/Zygote/FPUm3/src/lib/broadcast.jl:187 [inlined]
  [6] _pullback(__context__::Zygote.Context, 714::typeof(Base.Broadcast.broadcasted), 715::SparseArrays.HigherOrderFns.SparseVecStyle, f::Type{Float32}, args::SparseVector{Float64, Int64})
    @ Zygote ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:65

But perhaps Zygote ought to have a rule for broadcasting (::Type).(x), since what's written here would work for arbitrary arrays I believe, not just sparse, and would be faster than what it does now.

[CarloLucibello]

maybe it's ok to have the first broadcast here covering a limited use case not currently supported by Zygote and see how it goes? There are plenty of issues related to sparse arrays in the ML ecosystem, I'd like to get things going for downstream packages

[mzgubic]

Would merging the findnz and separating broadcasting out in a separate PR be an option? Since findnz is complete and broadcasting possibly requires in-depth discussion?

[CarloLucibello]

ok

[CarloLucibello]

The non-controversial findnz part is now in #590. After merging that, I will update this PR to contain the broadcast part only

[CarloLucibello]

So this is just a very specific instance of the general rule for broadcasting we have in Zygote. Would it be ok to merge this now and go with some generic rule (the one from Zygote or something else) later?

[mzgubic]

Should we close this given FluxML/Zygote.jl#1171 or would you prefer to keep it open?

[CarloLucibello]

Yeah, this is not needed anymore.