=Bring over OneElement for scalar getindex

Author: oxinabox

src/rulesets/Base/indexing.jl

@@ -81,16 +81,37 @@ 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...)
-    return ProjectTo(x)(dx)  # since we have x, may as well do this inside, not in rules
+    if plain_inds isa NTuple{N, Int} && T<:Number
+        # scalar indexing
+        return OneElement(dy, plain_inds, axes(x))
+    else  # some from slicing (potentially noncontigous)
+        dx = _setindex_zero(x, dy, plain_inds...)
+        ∇getindex!(dx, dy, plain_inds...)
+        return ProjectTo(x)(dx)  # since we have x, may as well do this inside, not in rules
+    end
 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))
+# TODO: should we teach ProjectTo that OneElement is more structurally sparse than anything it intersects nonstructural zeros with?
+
 """
     _setindex_zero(x, dy, inds...)
 

test/rulesets/Base/array.jl

@@ -398,8 +398,8 @@ end
     @test res == @inferred unthunk(rrule(imum, [1,2,1,2,1,2])[2](1.0)[2])
 
     # Structured matrix -- NB the minimum is a structral zero here
-    @test unthunk(rrule(imum, Diagonal(rand(3) .+ 1))[2](5.5)[2]) isa Diagonal
-    @test unthunk(rrule(imum, UpperTriangular(rand(3,3) .+ 1))[2](5.5)[2]) isa UpperTriangular{Float64}
+    @test unthunk(rrule(imum, Diagonal(rand(3) .+ 1))[2](5.5)[2]) isa Union{Diagonal, OneElement}
+    @test unthunk(rrule(imum, UpperTriangular(rand(3,3) .+ 1))[2](5.5)[2]) isa Union{UpperTriangular{Float64}, ChainRules.OneElement{Float64}} # must be at least as structured
     @test_skip test_rrule(imum, Diagonal(rand(3) .+ 1)) # MethodError: no method matching zero(::Type{Any}), from fill!(A::SparseArrays.SparseMatrixCSC{Any, Int64}, x::Bool)
 end
 

test/rulesets/Base/indexing.jl

@@ -89,7 +89,7 @@
 
         test_rrule(getindex, Symmetric(rand(3, 3)), 2, 2)
         sgrad = rrule(getindex, Symmetric(rand(3, 3)), 2, 3)[2](1.0)[2]
-        @test unthunk(sgrad) ≈ [0 0 0; 0 0 1/2; 0 1/2 0]
+        @test unthunk(sgrad) ≈ [0 0 0; 0 0 1/2; 0 1/2 0]  # We are actually getting this wrong now
     end
 
     @testset "getindex(::Array{<:Array})" begin