Merge branch 'main' of https://github.com/JuliaDiff/ChainRules.jl into sparsedet

Author: ElOceanografo

.github/workflows/CI.yml

@@ -21,7 +21,7 @@ jobs:
           - x86
           - x64
     steps:
-      - uses: actions/checkout@v3.5.2
+      - uses: actions/checkout@v3.5.3
       - uses: julia-actions/setup-julia@v1
         with:
           version: ${{ matrix.version }}

.github/workflows/IntegrationTest.yml

@@ -25,14 +25,14 @@ jobs:
         #    package: {user: JuliaDiff, repo: Diffractor.jl}
 
     steps:
-      - uses: actions/checkout@v3.5.2
+      - uses: actions/checkout@v3.5.3
       - uses: julia-actions/setup-julia@v1
         with:
           version: ${{ matrix.julia-version }}
           arch: x64
       - uses: julia-actions/julia-buildpkg@latest
       - name: Clone Downstream
-        uses: actions/checkout@v3.5.2
+        uses: actions/checkout@v3.5.3
         with:
           repository: ${{ matrix.package.user }}/${{ matrix.package.repo }}
           path: downstream

.github/workflows/JuliaNightly.yml

@@ -23,7 +23,7 @@ jobs:
           - x86
           - x64
     steps:
-      - uses: actions/checkout@v3.5.2
+      - uses: actions/checkout@v3.5.3
       - uses: julia-actions/setup-julia@v1
         with:
           version: ${{ matrix.version }}

.github/workflows/VersionVigilante_pull_request.yml

@@ -6,7 +6,7 @@ jobs:
   VersionVigilante:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v3.5.2
+      - uses: actions/checkout@v3.5.3
       - uses: julia-actions/setup-julia@latest
       - name: VersionVigilante.main
         id: versionvigilante_main

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRules"
 uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
-version = "1.49.0"
+version = "1.53.0"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/rulesets/Base/arraymath.jl

@@ -210,7 +210,13 @@ end # VERSION
 ##### `muladd`
 #####
 
-function frule((_, ΔA, ΔB, Δz), ::typeof(muladd), A, B, z)
+function frule(
+    (_, ΔA, ΔB, Δz),
+    ::typeof(muladd),
+    A::AbstractVecOrMat{<:CommutativeMulNumber},
+    B::AbstractVecOrMat{<:CommutativeMulNumber},
+    z::Union{CommutativeMulNumber, AbstractVecOrMat{<:CommutativeMulNumber}}
+)
     Ω = muladd(A, B, z)
     return Ω, ΔA * B .+ A * ΔB .+ Δz
 end
@@ -342,20 +348,43 @@ function rrule(::typeof(\), A::AbstractVecOrMat{<:Real}, B::AbstractVecOrMat{<:R
     project_B = ProjectTo(B)
 
     Y = A \ B
+
     function backslash_pullback(ȳ)
         Ȳ = unthunk(ȳ)
+        
+        Ȳf = Ȳ
+        @static if VERSION >= v"1.9"
+            # Need to ensure Ȳ is an array since since https://github.com/JuliaLang/julia/pull/44358
+            if !isa(Ȳ, AbstractArray)
+                Ȳf = [Ȳ]
+            end
+        end
+        Yf = Y
+        @static if VERSION >= v"1.9" 
+            # Need to ensure Yf is an array since since https://github.com/JuliaLang/julia/pull/44358
+            if !isa(Y, AbstractArray)
+                Yf = [Y]
+            end
+        end
+        #@info "vars" typeof(Ȳ) typeof(Y) typeof(Yf) typeof(A) typeof(B)
         ∂A = @thunk begin
-            B̄ = A' \ Ȳ
+            B̄ = A' \ Ȳf
             Ā = -B̄ * Y'
-            Ā = add!!(Ā, (B - A * Y) * B̄' / A')
-            Ā = add!!(Ā, A' \ Y * (Ȳ' - B̄'A))
+            t = (B - A * Y) * B̄'
+            @static if VERSION >= v"1.9" 
+                # Need to ensure t is an array since since https://github.com/JuliaLang/julia/pull/44358
+                if !isa(t, AbstractArray)
+                    t = [t]
+                end
+            end
+            Ā = add!!(Ā, t / A')
+            Ā = add!!(Ā, A' \ Yf * (Ȳ' - B̄'A))
             project_A(Ā)
         end
-        ∂B = @thunk project_B(A' \ Ȳ)
+        ∂B = @thunk project_B(A' \ Ȳf)
         return NoTangent(), ∂A, ∂B
     end
     return Y, backslash_pullback
-
 end
 
 #####

src/rulesets/Base/base.jl

@@ -3,8 +3,8 @@
 
 @scalar_rule copysign(y, x) (ifelse(signbit(x)!=signbit(y), -one(y), +one(y)), NoTangent())
 
-@scalar_rule one(x) zero(x)
-@scalar_rule zero(x) zero(x)
+@scalar_rule one(x) ZeroTangent()
+@scalar_rule zero(x) ZeroTangent()
 @scalar_rule transpose(x) true
 
 # `adjoint`

src/rulesets/Base/fastmath_able.jl

@@ -169,8 +169,8 @@ let
         @scalar_rule x / y (one(x) / y, -(Ω / y))
 
         ## many-arg +
-        function frule((_, Δx, Δy...), ::typeof(+), x::Number, ys::Number...)
-            +(x, ys...), +(Δx, Δy...)
+        function frule(Δs, ::typeof(+), x::Number, ys::Number...)
+            +(x, ys...), +(Base.tail(Δs)...)
         end
 
         function rrule(::typeof(+), x::Number, ys::Number...)

src/rulesets/Base/indexing.jl

@@ -1,14 +1,6 @@
-#####
-##### getindex(::Tuple)
-#####
-
-function frule((_, ẋ), ::typeof(getindex), x::Tuple, i::Integer)
-    return x[i], ẋ[i]
-end
-
-function frule((_, ẋ), ::typeof(getindex), x::Tuple, i)
-    y = x[i]
-    return y, Tangent{typeof(y)}(ẋ[i]...)
+# Int rather than Int64/Integer is intentional
+function frule((_, ẋ), ::typeof(getfield), x::Tuple, i::Int)
+    return x.i, ẋ.i
 end
 
 "for a given tuple type, returns a Val{N} where N is the length of the tuple"
@@ -77,20 +69,52 @@ end
 """
     ∇getindex(x, dy, inds...)
 
-For the `rrule` of `y = x[inds...]`, this function is roughly 
+For the `rrule` of `y = x[inds...]`, this function is roughly
 `setindex(zero(x), dy, inds...)`, returning the array `dx`.
 Differentiable. Includes `ProjectTo(x)(dx)`.
 """
-function ∇getindex(x::AbstractArray, dy, inds...)
+function ∇getindex(x::AbstractArray{T,N}, dy, inds...) where {T,N}
     # `to_indices` removes any logical indexing, colons, CartesianIndex etc,
     # leaving just Int / AbstractVector of Int
     plain_inds = Base.to_indices(x, inds)
-    dx = _setindex_zero(x, dy, plain_inds...)
-    ∇getindex!(dx, dy, plain_inds...)
+    dx = if plain_inds isa NTuple{N, Int} && T<:Number
+        # scalar indexing
+        OneElement(dy, plain_inds, axes(x))
+    else  # some from slicing (potentially noncontigous)
+        dx = _setindex_zero(x, dy, plain_inds...)
+        ∇getindex!(dx, dy, plain_inds...)
+    end
     return ProjectTo(x)(dx)  # since we have x, may as well do this inside, not in rules
 end
 ∇getindex(x::AbstractArray, z::AbstractZero, inds...) = z
 
+"""
+    OneElement(val, ind, axes) <: AbstractArray
+
+Extremely simple `struct` used for the gradient of scalar `getindex`.
+"""
+struct OneElement{T,N,I,A} <: AbstractArray{T,N}
+  val::T
+  ind::I
+  axes::A
+  OneElement(val::T, ind::I, axes::A) where {T<:Number, I<:NTuple{N,Int}, A<:NTuple{N,AbstractUnitRange}} where {N} = new{T,N,I,A}(val, ind, axes)
+end
+Base.size(A::OneElement) = map(length, A.axes)
+Base.axes(A::OneElement) = A.axes
+Base.getindex(A::OneElement{T,N}, i::Vararg{Int,N}) where {T,N} = ifelse(i==A.ind, A.val, zero(T))
+
+function ChainRulesCore.add!!(xs::AbstractArray{<:Any,N}, oe::OneElement{<:Any,N}) where {N}
+    if !ChainRulesCore.is_inplaceable_destination(xs)
+        xs = collect(xs)
+    end
+    xs[oe.ind...] += oe.val
+    return xs
+end
+
+Base.:(+)(xs::AbstractArray, oe::OneElement) = add!!(copy(xs), oe)
+Base.:(+)(oe::OneElement, xs::AbstractArray) = +(xs, oe)
+Base.:(+)(oe1::OneElement, oe2::OneElement) = +(collect(oe1), oe2)
+
 """
     _setindex_zero(x, dy, inds...)
 
@@ -159,29 +183,6 @@ function ∇getindex!(dx::AbstractGPUArray, dy, inds...)
     return dx
 end
 
-#####
-##### first, tail
-#####
-
-function frule((_, ẋ), ::typeof(first), x::Tuple)
-    return first(x), first(ẋ)
-end
-
-function rrule(::typeof(first), x::T) where {T<:Tuple}
-    first_back(dy) = (NoTangent(), Tangent{T}(ntuple(j -> j == 1 ? dy : NoTangent(), _tuple_N(T))...))
-    return first(x), first_back
-end
-
-function frule((_, ẋ), ::typeof(Base.tail), x::Tuple)
-    y = Base.tail(x)
-    return y, Tangent{typeof(y)}(Base.tail(ẋ)...)
-end
-
-function rrule(::typeof(Base.tail), x::T) where {T<:Tuple}
-    tail_pullback(dy) = (NoTangent(), Tangent{T}(NoTangent(), dy...))
-    return Base.tail(x), tail_pullback
-end
-
 #####
 ##### view
 #####

test/rulesets/Base/array.jl

@@ -358,14 +358,15 @@ end
     @test_skip test_frule(findmin, rand(3,4), output_tangent = (rand(), NoTangent()))
     @test_skip test_frule(findmin, rand(3,4), fkwargs=(dims=1,))
     # These skipped tests might be fixed by https://github.com/JuliaDiff/FiniteDifferences.jl/issues/188
+    # or by https://github.com/JuliaLang/julia/pull/48404
 
     # Reverse
     test_rrule(findmin, rand(10), output_tangent = (rand(), false))
     test_rrule(findmax, rand(10), output_tangent = (rand(), false))
-    test_rrule(findmin, rand(5,3))
-    test_rrule(findmax, rand(5,3))
-    @test [0 0; 0 5] == @inferred unthunk(rrule(findmax, [1 2; 3 4])[2]((5.0, nothing))[2])
-    @test [0 0; 0 5] == @inferred unthunk(rrule(findmax, [1 2; 3 4])[2]((5.0, NoTangent()))[2])
+    test_rrule(findmin, rand(5,3); check_inferred=false)
+    test_rrule(findmax, rand(5,3); check_inferred=false)
+    @test [0 0; 0 5] == unthunk(rrule(findmax, [1 2; 3 4])[2]((5.0, nothing))[2])
+    @test [0 0; 0 5] == unthunk(rrule(findmax, [1 2; 3 4])[2]((5.0, NoTangent()))[2])
 
     # Reverse with dims:
     @test [0 0; 5 6] == @inferred unthunk(rrule(findmax, [1 2; 3 4], dims=1)[2](([5 6], nothing))[2])
@@ -385,7 +386,7 @@ end
 
     # Reverse
     test_rrule(imum, rand(10))
-    test_rrule(imum, rand(3,4))
+    test_rrule(imum, rand(3,4); check_inferred=false)
     @gpu test_rrule(imum, rand(3,4), fkwargs=(dims=1,))
     test_rrule(imum, rand(3,4,5), fkwargs=(dims=(1,3),))