add tests for Iterators.take

src/lib/array.jl

@@ -323,9 +323,10 @@ end
 takefunc(itr, dy) = _restore(dy, _tryaxes(itr))
 
 @adjoint function Iterators.take(itr, n)
+  take_pullback(::AbstractArray{Nothing}) = nothing
   take_pullback(dy::NamedTuple{(:xs,:n)}) = (dy.xs, dy.n)
   take_pullback(dy::NamedTuple{(:n,:xs)}) = (dy.xs, dy.n)
-  take_pullback(dy) = (takefunc(itr, dy), nothing)
+  take_pullback(dy::AbstractArray) = (takefunc(itr, dy), nothing)
   Iterators.take(itr, n), take_pullback
 end
 

test/lib/array.jl

@@ -50,6 +50,23 @@ end
     @test @inferred back(collect(y)) == (nothing, [1.0, 0.0, 0.0, 0.0, 0.0], fill(1.0))
 end
 
+@testset "adjoints of Iterators.take" begin
+    y, back = _pullback(Iterators.take, 1:5, 3)
+    @test back(collect(y)) == (nothing, [1.0, 2.0, 3.0, 0.0, 0.0], nothing)
+    @test back([nothing for i in 1:3]) === nothing
+
+    @test gradient(x -> sum([2y for y in Iterators.take(x, 4)]), [1,2,3,4])[1] ≈ [2, 2, 2, 2]
+    @test gradient(x -> sum(2y for y in Iterators.take(x, 4)), [1,2,3,4])[1] ≈ [2, 2, 2, 2]
+
+    for p in (1.0, fill(1.0), [1.0])
+        @test gradient(p_ -> sum(map(prod, Iterators.take(p_, 1))), p) == (p,)
+        @test gradient(p_ -> sum(x for x in Iterators.take(p_, 1)), p) == (p,)
+    end
+
+    y, back = _pullback(Iterators.take, ones(2, 2), 3)
+    @test @inferred back(collect(y)) == (nothing, [1.0 1.0; 1.0 0.0], nothing)
+end
+
 @testset "collect" begin
     @testset "Dict" begin
         d = Dict(1 => 5, 2 => 6)
@@ -97,6 +114,16 @@ end
         @test gradient(x -> sum(broadcast(prod, Iterators.zip(x,x.^2))), ones(4)) == (3ones(4),)
         @test gradient(x -> sum(broadcast(prod, Iterators.zip(x.^2,x.^2))), ones(4)) == (4ones(4),)
     end
+
+
+    @testset "Iterators.Take" begin
+        z = Iterators.take(1:3, 2)
+        g = gradient(z -> sum(collect(z)), z)[1]
+        @test g == (xs=[1.0, 1.0, 0.0], n=nothing)
+
+        @test gradient(x -> sum(broadcast(prod, Iterators.take(x,2))), ones(4)) == ([1.0,1.0,0.0,0.0],)
+        @test gradient(x -> sum(broadcast(prod, Iterators.take(x.^2,2))), ones(4)) == (2*[1.0,1.0,0.0,0.0],)
+    end
 end
 
 @testset "dictionary comprehension" begin