use Base definition of stack

Author: CarloLucibello

Project.toml

@@ -5,6 +5,7 @@ version = "0.3.0"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 DataAPI = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a"
 DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab"
 FLoops = "cc61a311-1640-44b5-9fba-1b764f453329"

src/MLUtils.jl

@@ -22,7 +22,11 @@ import NNlib
 @traitdef IsTable{X}
 @traitimpl IsTable{X} <- Tables.istable(X)
 
-
+if VERSION < v"1.9.0-DEV.1163"
+    import Compat: stack
+else
+    import Base: stack
+end
 include("observation.jl")
 export numobs,
        getobs,
@@ -75,7 +79,7 @@ export batch,
        rand_like,
        randn_like,
        rpad_constant,
-       stack,
+       stack, # in Base since julia v1.9
        unbatch,
        unsqueeze,
        unstack,

src/deprecations.jl

@@ -1,5 +1,4 @@
 # Deprecated in v0.2
-@deprecate stack(x, dims) stack(x; dims=dims)
 @deprecate unstack(x, dims) unstack(x; dims=dims)
 @deprecate unsqueeze(x::AbstractArray, dims::Int) unsqueeze(x; dims=dims)
 @deprecate unsqueeze(dims::Int) unsqueeze(dims=dims)

src/utils.jl

@@ -57,72 +57,6 @@ _unsqueeze(x, dims) = unsqueeze(x; dims)
 
 Base.show_function(io::IO, u::Base.Fix2{typeof(_unsqueeze)}, ::Bool) = print(io, "unsqueeze(dims=", u.x, ")")
 
-"""
-    stack(xs; dims)
-
-Concatenate the given array of arrays `xs` into a single array along the
-new dimension `dims`. All arrays need to be of the same size.
-
-See also [`unsqueeze`](@ref), [`unstack`](@ref) and [`batch`](@ref).
-
-# Examples
-
-```jldoctest
-julia> xs = [[1, 2], [3, 4], [5, 6]]
-3-element Vector{Vector{Int64}}:
- [1, 2]
- [3, 4]
- [5, 6]
-
-julia> stack(xs, dims=1)
-3×2 Matrix{Int64}:
- 1  2
- 3  4
- 5  6
-
-julia> stack(xs, dims=2)
-2×3 Matrix{Int64}:
- 1  3  5
- 2  4  6
-
-julia> stack(xs, dims=3)
-2×1×3 Array{Int64, 3}:
-[:, :, 1] =
- 1
- 2
-
-[:, :, 2] =
- 3
- 4
-
-[:, :, 3] =
- 5
- 6
-```
-"""
-function stack(xs; dims::Int)
-    N = ndims(xs[1])
-    if dims <= N
-        vs = unsqueeze.(xs; dims)
-    else
-        vs = xs
-    end
-    if dims == 1
-        return reduce(vcat, vs)
-    elseif dims === 2
-        return reduce(hcat, vs)
-    else
-        return reduce((x, y) -> cat(x, y; dims=dims), vs)
-    end
-end
-
-function rrule(::typeof(stack), xs; dims::Int)
-    function stack_pullback(Δ)
-        return (NoTangent(), unstack(unthunk(Δ); dims=dims))
-    end
-    return stack(xs; dims=dims), stack_pullback
-end
-
 """
     unstack(xs; dims)
 

test/utils.jl

@@ -14,7 +14,7 @@ end
     x = randn(3,3)
     stacked = stack([x, x], dims=2)
     @test size(stacked) == (3,2,3)
-    @test_broken @inferred(stack([x, x], dims=2)) == stacked
+    @test @inferred(stack([x, x], dims=2)) == stacked
 
     stacked_array=[ 8 9 3 5; 9 6 6 9; 9 1 7 2; 7 4 10 6 ]
     unstacked_array=[[8, 9, 9, 7], [9, 6, 1, 4], [3, 6, 7, 10], [5, 9, 2, 6]]
@@ -30,7 +30,6 @@ end
     a = [[1] for i in 1:10000]
     @test size(stack(a, dims=1)) == (10000, 1)
     @test size(stack(a, dims=2)) == (1, 10000)
-    @test size(stack(a, dims=3)) == (1, 1, 10000)
 end
 
 @testset "batch and unbatch" begin