start working on spdiag rrule

Author: ElOceanografo

Project.toml

@@ -5,6 +5,7 @@ version = "1.49.0"
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"

src/rulesets/SparseArrays/sparsematrix.jl

@@ -49,3 +49,33 @@ function rrule(::typeof(findnz), v::AbstractSparseVector)
 
     return (I, V), findnz_pullback
 end
+
+function _spdiagm_back(p, ȳ)
+    k, v = p
+    d = diag(unthunk(ȳ), k)[1:length(v)] # handle if diagonal was smaller than matrix
+    return Tangent{typeof(p)}(second = d)
+end
+
+function rrule(::typeof(spdiagm), m::Integer, n::Integer, kv::Pair{<:Integer,<:AbstractVector}...)
+    function diagm_pullback(Δ)
+        _, ȳ = unthunk(Δ)
+        return (NoTangent(), NoTangent(), NoTangent(), _spdiagm_back.(kv, Ref(ȳ))...)
+    end
+    return spdiagm(m, n, kv...), diagm_pullback
+end
+
+function rrule(::typeof(spdiagm), kv::Pair{<:Integer,<:AbstractVector}...)
+    function diagm_pullback(Δ)
+        _, ȳ = unthunk(Δ)
+        return (NoTangent(), _spdiagm_back.(kv, Ref(ȳ))...)
+    end
+    return spdiagm(kv...), diagm_pullback
+end
+
+function rrule(::typeof(spdiagm), v::AbstractVector)
+    function diagm_pullback(Δ)
+        _, ȳ = unthunk(Δ)
+        return (NoTangent(), diag(ȳ))
+    end
+    return spdiagm(v), diagm_pullback
+end

test/rulesets/SparseArrays/sparsematrix.jl

@@ -18,6 +18,18 @@ end
     test_rrule(SparseVector{Float32}, Float32.(v), rtol=1e-4)
 end
 
+@testset "spdiagm" begin
+    @test 1 == 1
+    m = 5
+    n = 4
+    v1 = ones(m)
+    v2 = ones(n)
+    test_rrule(spdiagm, m, n, 0 => v2)
+
+    # test_rrule(spdiagm, 0 => v1)
+    # test_rrule(spdiagm, v1)
+end
+
 @testset "findnz" begin
     A = sprand(5, 5, 0.5)
     dA = similar(A)

test/runtests.jl

@@ -52,41 +52,41 @@ end
 
     test_method_tables()  # Check the global method tables are consistent
 
-    # Each file puts all tests inside one or more @testset blocks
-    include_test("rulesets/Base/base.jl")
-    include_test("rulesets/Base/fastmath_able.jl")
-    include_test("rulesets/Base/evalpoly.jl")
-    include_test("rulesets/Base/array.jl")
-    include_test("rulesets/Base/arraymath.jl")
-    include_test("rulesets/Base/indexing.jl")
-    include_test("rulesets/Base/mapreduce.jl")
-    include_test("rulesets/Base/sort.jl")
-    include_test("rulesets/Base/broadcast.jl")
-
-    include_test("unzipped.jl")  # used primarily for broadcast
+    # # Each file puts all tests inside one or more @testset blocks
+    # include_test("rulesets/Base/base.jl")
+    # include_test("rulesets/Base/fastmath_able.jl")
+    # include_test("rulesets/Base/evalpoly.jl")
+    # include_test("rulesets/Base/array.jl")
+    # include_test("rulesets/Base/arraymath.jl")
+    # include_test("rulesets/Base/indexing.jl")
+    # include_test("rulesets/Base/mapreduce.jl")
+    # include_test("rulesets/Base/sort.jl")
+    # include_test("rulesets/Base/broadcast.jl")
+
+    # include_test("unzipped.jl")  # used primarily for broadcast
+
+    # println()
+
+    # include_test("rulesets/Statistics/statistics.jl")
+
+    # println()
+
+    # include_test("rulesets/LinearAlgebra/dense.jl")
+    # include_test("rulesets/LinearAlgebra/norm.jl")
+    # include_test("rulesets/LinearAlgebra/matfun.jl")
+    # include_test("rulesets/LinearAlgebra/structured.jl")
+    # include_test("rulesets/LinearAlgebra/symmetric.jl")
+    # include_test("rulesets/LinearAlgebra/factorization.jl")
+    # include_test("rulesets/LinearAlgebra/blas.jl")
+    # include_test("rulesets/LinearAlgebra/lapack.jl")
+    # include_test("rulesets/LinearAlgebra/uniformscaling.jl")
 
     println()
 
-    include_test("rulesets/Statistics/statistics.jl")
+    include("rulesets/SparseArrays/sparsematrix.jl")
 
     println()
 
-    include_test("rulesets/LinearAlgebra/dense.jl")
-    include_test("rulesets/LinearAlgebra/norm.jl")
-    include_test("rulesets/LinearAlgebra/matfun.jl")
-    include_test("rulesets/LinearAlgebra/structured.jl")
-    include_test("rulesets/LinearAlgebra/symmetric.jl")
-    include_test("rulesets/LinearAlgebra/factorization.jl")
-    include_test("rulesets/LinearAlgebra/blas.jl")
-    include_test("rulesets/LinearAlgebra/lapack.jl")
-    include_test("rulesets/LinearAlgebra/uniformscaling.jl")
-
-    println()
-
-    include_test("rulesets/SparseArrays/sparsematrix.jl")
-
-    println()
-
-    include_test("rulesets/Random/random.jl")
+    # include_test("rulesets/Random/random.jl")
     println()
 end