Use CR OneElement

Project.toml

@@ -28,7 +28,7 @@ ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
 [compat]
 AbstractFFTs = "1.3.1"
-ChainRules = "1.44.1"
+ChainRules = "1.51.0"
 ChainRulesCore = "1.9"
 ChainRulesTestUtils = "1"
 Colors = "0.12"

src/Zygote.jl

@@ -7,7 +7,7 @@ import ZygoteRules: @adjoint, @adjoint!, AContext, adjoint, _pullback, pullback,
   literal_getproperty, literal_getfield, unthunk_tangent
 
 using ChainRulesCore
-using ChainRules: ChainRules, rrule, unthunk, canonicalize
+using ChainRules: ChainRules, rrule, unthunk, canonicalize, OneElement
 using IRTools
 using MacroTools, Requires
 using MacroTools: @forward

src/lib/array.jl

@@ -41,21 +41,6 @@ end
 @adjoint (::Type{T})(sz) where {T<:Zeros} = T(sz), Δ->(nothing,)
 @adjoint (::Type{T})(sz) where {T<:Ones} = T(sz), Δ->(nothing,)
 
-"""
-    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))
-
 
 _zero(xs::AbstractArray{<:Number}, T::Type{Nothing}) = fill!(similar(xs), zero(eltype(xs)))
 _zero(xs::AbstractArray{<:Number}, T) = fill!(similar(xs, T), false)