Implement reverse and getindex for AbstractZero and thunks (#522)

Author: mzgubic

src/tangent_types/abstract_zero.jl

@@ -30,11 +30,12 @@ Base.convert(::Type{T}, x::AbstractZero) where {T<:Number} = zero(T)
 (::Type{Complex})(x::AbstractZero, y::Real) = Complex(false, y)
 (::Type{Complex})(x::Real, y::AbstractZero) = Complex(x, false)
 
-Base.getindex(z::AbstractZero, k) = z
+Base.getindex(z::AbstractZero, args...) = z
 
 Base.view(z::AbstractZero, ind...) = z
 Base.sum(z::AbstractZero; dims=:) = z
 Base.reshape(z::AbstractZero, size...) = z
+Base.reverse(z::AbstractZero, args...; kwargs...) = z
 
 """
     ZeroTangent() <: AbstractZero

src/tangent_types/thunks.jl

@@ -47,6 +47,7 @@ Base.sum!(r, A::AbstractThunk; kws...) = sum!(r, unthunk(A); kws...)
 Base.fill(a::AbstractThunk, b::Integer) = fill(unthunk(a), b)
 Base.vec(a::AbstractThunk) = vec(unthunk(a))
 Base.reshape(a::AbstractThunk, args...) = reshape(unthunk(a), args...)
+Base.reverse(a::AbstractThunk, args...; kwargs...) = reverse(unthunk(a), args...; kwargs...)
 Base.getindex(a::AbstractThunk, args...) = getindex(unthunk(a), args...)
 Base.setindex!(a::AbstractThunk, value, key...) = throw(MutateThunkException())
 Base.selectdim(a::AbstractThunk, args...) = selectdim(unthunk(a), args...)

test/tangent_types/abstract_zero.jl

@@ -13,6 +13,11 @@
 
         @test reshape(NoTangent(), (1, :)) === NoTangent()
         @test reshape(ZeroTangent(), 2, 3, 4, 5) === ZeroTangent()
+
+        @test reverse(NoTangent()) === NoTangent()
+        @test reverse(ZeroTangent()) === ZeroTangent()
+        @test reverse(ZeroTangent(); dims=1) === ZeroTangent()
+        @test reverse(ZeroTangent(), 2, 5) === ZeroTangent()
     end
 
     @testset "ZeroTangent" begin
@@ -76,6 +81,11 @@
         @test convert(Int64, ZeroTangent()) === Int64(0)
         @test convert(Float32, ZeroTangent()) === 0.0f0
         @test convert(ComplexF64, ZeroTangent()) === 0.0 + 0.0im
+
+        @test z[1] === z
+        @test z[1:3] === z
+        @test z[1, 2] === z
+        @test getindex(z) === z
     end
 
     @testset "NoTangent" begin
@@ -115,6 +125,11 @@
 
         @test convert(Int64, NoTangent()) == 0
         @test convert(Float64, NoTangent()) == 0.0
+
+        @test dne[1] === dne
+        @test dne[1:3] === dne
+        @test dne[1, 2] === dne
+        @test getindex(dne) === dne
     end
 
     @testset "ambiguities" begin

test/tangent_types/thunks.jl

@@ -107,6 +107,9 @@
         v = [1, 2, 3]
         t = @thunk(v)
 
+        m = rand(3, 3)
+        tm = @thunk(m)
+
         if VERSION >= v"1.2"
             @test 3 == mapreduce(_ -> 1, +, t)
             @test 3 == mapreduce((_, _) -> 1, +, v, t)
@@ -120,6 +123,10 @@
         @test 1 == getindex(t, 1)
         @test_throws MutateThunkException setindex!(t, 0.0, 1)
         @test [4; 5; 6] == selectdim([1 2 3; 4 5 6], 1, 2)
+
+        @test reverse(t) == reverse(v)
+        @test reverse(t, 2) == reverse(v, 2)
+        @test reverse(tm; dims=2) == reverse(m; dims=2)
     end
 
     @testset "LinearAlgebra" begin