Add TestIterator for testing rules with iterator primal inputs (#54)

* Implement test iterator

* Export iterator

* Test iterator

* Test testers on iterator

* Increment version number

* Implement and test hash

Author: sethaxen

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRulesTestUtils"
 uuid = "cdddcdb0-9152-4a09-a978-84456f9df70a"
-version = "0.5.0"
+version = "0.5.1"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/ChainRulesTestUtils.jl

@@ -11,11 +11,13 @@ using Test
 
 const _fdm = central_fdm(5, 1)
 
+export TestIterator
 export test_scalar, frule_test, rrule_test, generate_well_conditioned_matrix
 
 include("generate_tangent.jl")
 include("to_vec.jl")
 include("isapprox.jl")
 include("data_generation.jl")
+include("iterator.jl")
 include("testers.jl")
 end # module

src/iterator.jl

@@ -0,0 +1,66 @@
+"""
+    TestIterator{T,IS<:Base.IteratorSize,IE<:Base.IteratorEltype}
+
+A configurable iterator for testing purposes.
+
+    TestIterator(data, itersize, itereltype)
+    TestIterator(data)
+
+The iterator wraps another iterator `data`, such as an array, that must have at least as
+many features implemented as the test iterator and have a `FiniteDifferences.to_vec`
+overload. By default, the iterator it has the same features as `data`.
+
+The optional methods `eltype`, length`, and `size` are automatically defined and forwarded
+to `data` if the type arguments indicate that they should be defined.
+"""
+struct TestIterator{T,IS,IE}
+    data::T
+end
+function TestIterator(data, itersize::Base.IteratorSize, itereltype::Base.IteratorEltype)
+    return TestIterator{typeof(data),typeof(itersize),typeof(itereltype)}(data)
+end
+TestIterator(data) = TestIterator(data, Base.IteratorSize(data), Base.IteratorEltype(data))
+
+Base.iterate(iter::TestIterator) = iterate(iter.data)
+Base.iterate(iter::TestIterator, state) = iterate(iter.data, state)
+
+Base.IteratorSize(::Type{<:TestIterator{<:Any,IS}}) where {IS} = IS()
+
+Base.IteratorEltype(::Type{<:TestIterator{<:Any,<:Any,IE}}) where {IE} = IE()
+
+Base.eltype(::Type{<:TestIterator{T,<:Any,Base.HasEltype}}) where {T} = eltype(T)
+
+Base.length(iter::TestIterator{<:Any,Base.HasLength}) = length(iter.data)
+Base.length(iter::TestIterator{<:Any,<:Base.HasShape}) = length(iter.data)
+
+Base.size(iter::TestIterator{<:Any,<:Base.HasShape}) = size(iter.data)
+
+Base.:(==)(iter1::T, iter2::T) where {T<:TestIterator} = iter1.data == iter2.data
+
+Base.isequal(iter1::T, iter2::T) where {T<:TestIterator} = isequal(iter1.data, iter2.data)
+
+function Base.hash(iter::TestIterator{<:Any,IT,IS}) where {IT,IS}
+    return mapreduce(hash, hash, (iter.data, IT, IS))
+end
+
+Base.isapprox(iter1::TestIterator, iter2::TestIterator) = false
+function Base.isapprox(
+    iter1::TestIterator{T1,IS,IE},
+    iter2::TestIterator{T2,IS,IE};
+    kwargs...,
+) where {T1,T2,IS,IE}
+    return isapprox(iter1.data, iter2.data; kwargs...)
+end
+
+function rand_tangent(rng::AbstractRNG, x::TestIterator{<:Any,IS,IE}) where {IS,IE}
+    ∂data = rand_tangent(rng, x.data)
+    return TestIterator{typeof(∂data),IS,IE}(∂data)
+end
+
+function FiniteDifferences.to_vec(iter::TestIterator)
+    iter_vec, back = to_vec(iter.data)
+    function TestIterator_from_vec(v)
+        return TestIterator(back(v), Base.IteratorSize(iter), Base.IteratorEltype(iter))
+    end
+    return iter_vec, TestIterator_from_vec
+end

test/iterator.jl

@@ -0,0 +1,110 @@
+@testset "TestIterator" begin
+    @testset "Constructors" begin
+        data = randn(3)
+        iter = TestIterator(data)
+        @test iter isa TestIterator{
+            typeof(data),
+            typeof(Base.IteratorSize(data)),
+            typeof(Base.IteratorEltype(data)),
+        }
+        @test iter.data === data
+
+        data = randn(2, 3, 4)
+        iter = TestIterator(data)
+        @test iter isa TestIterator{
+            typeof(data),
+            typeof(Base.IteratorSize(data)),
+            typeof(Base.IteratorEltype(data)),
+        }
+        @test iter.data === data
+
+        data = randn(2, 3, 4)
+        iter = TestIterator(data, Base.SizeUnknown(), Base.EltypeUnknown())
+        @test iter isa TestIterator{typeof(data),Base.SizeUnknown,Base.EltypeUnknown}
+    end
+
+    @testset "iterate" begin
+        data = randn(3)
+        iter = TestIterator(data)
+
+        @test iterate(iter) === iterate(data)
+        _, state = iterate(data)
+        @test iterate(iter, state) === iterate(data, state)
+    end
+
+    @testset "optional interface methods" begin
+        data = randn(2, 3, 4)
+        iter = TestIterator(data)
+        @test eltype(iter) === eltype(data)
+        @test length(iter) === length(data)
+        @test size(iter) === size(data)
+
+        iter = TestIterator(data, Base.HasLength(), Base.HasEltype())
+        @test length(iter) === length(data)
+        @test_throws MethodError size(iter)
+        @test eltype(iter) === eltype(iter)
+
+        iter = TestIterator(data, Base.SizeUnknown(), Base.EltypeUnknown())
+        @test_throws MethodError length(iter)
+        @test eltype(iter) === Any
+    end
+
+    @testset "==" begin
+        data = randn(2, 3, 4)
+        iter1 = TestIterator(data, Base.HasLength(), Base.HasEltype())
+        iter2 = TestIterator(data, Base.HasLength(), Base.EltypeUnknown())
+        @test iter2 != iter1
+
+        iter3 = TestIterator(copy(data), Base.HasLength(), Base.HasEltype())
+        @test iter3 == iter1
+    end
+
+    @testset "isequal" begin
+        data = randn(2, 3, 4)
+        iter1 = TestIterator(data, Base.HasLength(), Base.HasEltype())
+        iter2 = TestIterator(data, Base.HasLength(), Base.EltypeUnknown())
+        @test !isequal(iter2, iter1)
+
+        iter3 = TestIterator(copy(data), Base.HasLength(), Base.HasEltype())
+        @test isequal(iter3, iter1)
+    end
+
+    @testset "hash" begin
+        data = randn(2, 3, 4)
+        iter1 = TestIterator(data, Base.HasLength(), Base.HasEltype())
+        iter2 = TestIterator(data, Base.HasLength(), Base.EltypeUnknown())
+        @test hash(iter2) != hash(iter1)
+
+        iter3 = TestIterator(copy(data), Base.HasLength(), Base.HasEltype())
+        @test hash(iter3) == hash(iter1)
+    end
+
+    @testset "isapprox" begin
+        data = randn(3)
+        iter1 = TestIterator(data, Base.HasLength(), Base.HasEltype())
+        iter2 = TestIterator(data, Base.HasLength(), Base.EltypeUnknown())
+        @test !isapprox(iter2, iter1)
+
+        iter3 = TestIterator(data .+ eps() .* rand.(), Base.HasLength(), Base.HasEltype())
+        @test isapprox(iter3, iter1)
+    end
+
+    @testset "to_vec" begin
+        data = randn(2, 3, 4)
+        iter = TestIterator(data, Base.SizeUnknown(), Base.EltypeUnknown())
+        v, back = ChainRulesTestUtils.to_vec(iter)
+        @test v isa AbstractVector{eltype(data)}
+        @test collect(v) == collect(vec(data))
+        iter2 = back(v)
+        @test iter2 == iter
+    end
+
+    @testset "rand_tangent" begin
+        data = randn(2, 3, 4)
+        iter = TestIterator(data, Base.SizeUnknown(), Base.EltypeUnknown())
+        ∂iter = rand_tangent(iter)
+        @test ∂iter isa typeof(iter)
+        @test size(∂iter.data) == size(iter.data)
+        @test eltype(∂iter.data) === eltype(iter.data)
+    end
+end

test/runtests.jl

@@ -8,6 +8,7 @@ using Test
     include("generate_tangent.jl")
     include("to_vec.jl")
     include("isapprox.jl")
+    include("iterator.jl")
     include("testers.jl")
     include("data_generation.jl")
 end

test/testers.jl

@@ -7,6 +7,39 @@ sinconj(x) = sin(x)
 
 primalapprox(x) = x
 
+function iterfun(iter)
+    state = iterate(iter)
+    state === nothing && error()
+    (x, i) = state
+    s = x^2
+    while true
+        state = iterate(iter, i)
+        state === nothing && break
+        (x, i) = state
+        s += x^2
+    end
+    return s
+end
+
+function ChainRulesCore.frule((_, Δiter), ::typeof(iterfun), iter)
+    iter_Δiter = zip(iter, Δiter)
+    state = iterate(iter_Δiter)
+    state === nothing && error()
+    # for some reason the following line errors if the frule is defined within a testset
+    ((x, Δx), i) = state
+    return iterfun(iter), sum(2 .* iter.data .* Δiter.data)
+    s = x^2
+    ∂s = 2 * x * Δx
+    while true
+        state = iterate(iter_Δiter, i)
+        state === nothing && break
+        ((x, Δx), i) = state
+        s += x^2
+        ∂s += 2 * x * Δx
+    end
+    return s, ∂s
+end
+
 @testset "testers.jl" begin
     @testset "test_scalar" begin
         double(x) = 2x
@@ -202,4 +235,28 @@ primalapprox(x) = x
         frule_test(primalapprox, (randn(), randn()); atol = 1e-6)
         rrule_test(primalapprox, randn(), (randn(), randn()); atol = 1e-6)
     end
+
+    @testset "TestIterator input" begin
+        function ChainRulesCore.rrule(::typeof(iterfun), iter::TestIterator)
+            function iterfun_pullback(Δs)
+                data = iter.data
+                ∂data = (2 * Δs) .* conj.(data)
+                ∂iter = TestIterator(
+                    ∂data,
+                    Base.IteratorSize(iter),
+                    Base.IteratorEltype(iter),
+                )
+                return (NO_FIELDS, ∂iter)
+            end
+            return iterfun(iter), iterfun_pullback
+        end
+
+        # define iterator with the minimal iterator interface
+        x = TestIterator(randn(2, 3), Base.SizeUnknown(), Base.EltypeUnknown())
+        ẋ = TestIterator(randn(2, 3), Base.SizeUnknown(), Base.EltypeUnknown())
+        x̄ = TestIterator(randn(2, 3), Base.SizeUnknown(), Base.EltypeUnknown())
+
+        frule_test(iterfun, (x, ẋ))
+        rrule_test(iterfun, randn(), (x, x̄))
+    end
 end