diff --git a/Project.toml b/Project.toml index 2a59656eb..a18c5c4b0 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "ChainRules" uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2" -version = "1.74.8" +version = "1.74.9" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" diff --git a/src/unzipped.jl b/src/unzipped.jl index fe5875e6f..1dc66c0bb 100644 --- a/src/unzipped.jl +++ b/src/unzipped.jl @@ -122,17 +122,27 @@ end """ unzip(t) -Also works on a tuple of tuples: +Also works on a non-empty tuple of tuples of equal length: ```jldoctest julia> unzip(((1,2), (30,40), (500,600))) ((1, 30, 500), (2, 40, 600)) ``` """ +unzip(xs::Tuple{Vararg{NTuple{N,Any}}}) where {N} = unzip(xs, Val(N)) +function unzip(::Tuple{}) + return throw( + ArgumentError("unzip cannot determine the number of components of an empty tuple") + ) +end function unzip(xs::Tuple) - x1 = first(xs) - x1 isa Tuple || throw(ArgumentError("unzip only accepts arrays or tuples of tuples")) - return ntuple(i -> map(Get(i), xs),length(x1)) + all(x -> x isa Tuple, xs) || + throw(ArgumentError("unzip only accepts arrays or tuples of tuples")) + return throw( + DimensionMismatch( + lazy"unzip requires tuples of equal length, got lengths $(map(length, xs))" + ), + ) end struct Get{i} end diff --git a/test/unzipped.jl b/test/unzipped.jl index 6bc508f3f..2d3cba4da 100644 --- a/test/unzipped.jl +++ b/test/unzipped.jl @@ -27,6 +27,20 @@ using ChainRules: unzip_broadcast, unzip #, unzip_map @test fun(tuple, (1,2,3), [4,5,6]) == ([1, 2, 3], [4, 5, 6]) # mix tuple & vector end + @testset "unzip tuple of tuples" begin + @test @inferred(unzip(((1, 2.0), (3, 4.0)))) == ((1, 3), (2.0, 4.0)) + @test @inferred(unzip(((1,), (2,)))) == ((1, 2),) + @test @inferred(unzip(((), ()))) == () + @test @inferred(unzip_broadcast(x -> (x, 2x), (1, 2.0))) == ((1, 2.0), (2, 4.0)) + + @test_throws ArgumentError unzip(()) + @test_throws ArgumentError unzip((1, (2, 3))) + @test_throws ArgumentError unzip(((1, 2), [3, 4])) + @test_throws ArgumentError unzip(((1, 2), 3)) + @test_throws DimensionMismatch unzip(((1, 2), (3,))) + @test_throws DimensionMismatch unzip(((1,), (2, 3))) + end + @testset "rrules" begin # These exist to allow for second derivatives