Some cat improvements (#52)

* improve cat functions

* handle reduce hcat fast cases

Author: mcabbott

src/array.jl

@@ -109,25 +109,32 @@ _onehot_bool_type(::OneHotLike{<:Any, <:Any, var"N+1", <:AbstractGPUArray}) wher
 _notall_onehot(x::OneHotArray, xs::OneHotArray...) = false
 _notall_onehot(x::OneHotLike, xs::OneHotLike...) = any(x -> !_isonehot(x), (x, xs...))
 
-function Base.cat(x::OneHotLike{<:Any, <:Any, N}, xs::OneHotLike...; dims::Int) where N
+function Base._cat(dims::Int, x::OneHotLike{<:Any, <:Any, N}, xs::OneHotLike...) where N
   if isone(dims) || _notall_onehot(x, xs...)
-    return cat(map(x -> convert(_onehot_bool_type(x), x), (x, xs...))...; dims = dims)
+    # return cat(map(x -> convert(_onehot_bool_type(x), x), (x, xs...))...; dims = dims)
+    return invoke(Base._cat, Tuple{Int, Vararg{AbstractArray{Bool}}}, dims, x, xs...)
   else
     L = _nlabels(x, xs...)
-
     return OneHotArray(cat(_indices(x), _indices.(xs)...; dims = dims - 1), L)
   end
 end
+function Base._cat(::Val{dims}, x::OneHotLike{<:Any, <:Any, N}, xs::OneHotLike...) where {N,dims}
+  if !(dims isa Integer) || isone(dims) || _notall_onehot(x, xs...)
+    # return cat(map(x -> convert(_onehot_bool_type(x), x), (x, xs...))...; dims = dims)
+    return invoke(Base._cat, Tuple{Val{dims}, Vararg{AbstractArray{Bool}}}, Val(dims), x, xs...)
+  else
+    L = _nlabels(x, xs...)
+    return OneHotArray(cat(_indices(x), _indices.(xs)...; dims = Val(dims - 1)), L)
+  end
+end
 
-Base.hcat(x::OneHotLike, xs::OneHotLike...) = cat(x, xs...; dims = 2)
-Base.vcat(x::OneHotLike, xs::OneHotLike...) =
-  vcat(map(x -> convert(_onehot_bool_type(x), x), (x, xs...))...)
+Base.hcat(x::OneHotLike, xs::OneHotLike...) = cat(x, xs...; dims = Val(2))
 
 # optimized concatenation for matrices and vectors of same parameters
 Base.hcat(x::OneHotMatrix, xs::OneHotMatrix...) =
-  OneHotMatrix(reduce(vcat, _indices.(xs); init = _indices(x)), _nlabels(x, xs...))
+  OneHotMatrix(vcat(_indices(x), _indices.(xs)...), _nlabels(x, xs...))
 Base.hcat(x::OneHotVector, xs::OneHotVector...) =
-  OneHotMatrix(reduce(vcat, _indices.(xs); init = _indices(x)), _nlabels(x, xs...))
+  OneHotMatrix(UInt32[_indices(x), _indices.(xs)...], _nlabels(x, xs...))
 
 if isdefined(Base, :stack)
   import Base: _stack
@@ -140,6 +147,13 @@ function _stack(::Colon, xs::AbstractArray{<:OneHotArray})
   OneHotArray(Compat.stack(_indices, xs), n)
 end
 
+Base.reduce(::typeof(hcat), xs::AbstractVector{<:OneHotArray{<:Any, 0, 1}}) = Compat.stack(xs)
+function Base.reduce(::typeof(hcat), xs::AbstractVector{<:OneHotMatrix})
+  n = _nlabels(first(xs))
+  all(x -> _nlabels(x)==n, xs) || throw(DimensionMismatch("The number of labels are not the same for all one-hot arrays."))
+  OneHotArray(reduce(vcat, _indices.(xs)), n)
+end
+
 Adapt.adapt_structure(T, x::OneHotArray) = OneHotArray(adapt(T, _indices(x)), x.nlabels)
 
 function Base.BroadcastStyle(::Type{<:OneHotArray{<:Any, <:Any, var"N+1", T}}) where {var"N+1", T <: AbstractGPUArray}

test/array.jl

@@ -52,12 +52,22 @@ end
   @test hcat(ov, ov) isa OneHotMatrix
   @test vcat(ov, ov) == vcat(collect(ov), collect(ov))
   @test cat(ov, ov; dims = 3) == OneHotArray(cat(ov.indices, ov.indices; dims = 2), 10)
+  @test cat(ov, ov; dims = 3) isa OneHotArray
+  @test cat(ov, ov; dims = Val(3)) == OneHotArray(cat(ov.indices, ov.indices; dims = 2), 10)
+  @test cat(ov, ov; dims = Val(3)) isa OneHotArray
+  @test cat(ov, ov; dims = (1, 2)) == cat(collect(ov), collect(ov); dims = (1, 2))
 
   # matrix cat
   @test hcat(om, om) == OneHotMatrix(vcat(om.indices, om.indices), 10)
   @test hcat(om, om) isa OneHotMatrix
-  @test vcat(om, om) == vcat(collect(om), collect(om))
+  @test vcat(om, om) == vcat(collect(om), collect(om))  # not one-hot!
+  @test cat(om, om; dims = 1) == vcat(collect(om), collect(om))
+  @test cat(om, om; dims = Val(1)) == vcat(collect(om), collect(om))
   @test cat(om, om; dims = 3) == OneHotArray(cat(om.indices, om.indices; dims = 2), 10)
+  @test cat(om, om; dims = 3) isa OneHotArray
+  @test cat(om, om; dims = Val(3)) == OneHotArray(cat(om.indices, om.indices; dims = 2), 10)
+  @test cat(om, om; dims = Val(3)) isa OneHotArray
+  @test cat(om, om; dims = (1, 2)) == cat(collect(om), collect(om); dims = (1, 2))
 
   # array cat
   @test cat(oa, oa; dims = 3) == OneHotArray(cat(oa.indices, oa.indices; dims = 2), 10)
@@ -71,9 +81,19 @@ end
   @test stack([om, om]) isa OneHotArray
   @test stack([oa, oa, oa, oa]) isa OneHotArray
 
+  # reduce(hcat)
+  @test reduce(hcat, [ov, ov]) == hcat(ov, ov)
+  @test reduce(hcat, [ov, ov]) isa OneHotMatrix
+  @test reduce(hcat, [onehotbatch(1, 1:3), onehotbatch(1, 1:3)]) == [1 1; 0 0; 0 0]
+  @test reduce(hcat, [onehotbatch(1, 1:3), onehotbatch(1, 1:3)]) isa OneHotMatrix
+  @test reduce(hcat, [om, om]) == hcat(om, om)
+  @test reduce(hcat, [om, om]) isa OneHotMatrix
+
   # proper error handling of inconsistent sizes
   @test_throws DimensionMismatch hcat(ov, ov2)
   @test_throws DimensionMismatch hcat(om, om2)
+  @test_throws DimensionMismatch stack([om, om2])
+  @test_throws DimensionMismatch reduce(hcat, [om, om2])
 end
 
 @testset "Base.reshape" begin
@@ -88,6 +108,8 @@ end
   @testset "w/ cat" begin
     r = reshape(oa, 10, :)
     @test hcat(r, r) isa OneHotArray
+    @test cat(r, r; dims = 2) isa OneHotArray
+    @test cat(r, r; dims = Val(2)) isa OneHotArray
     @test vcat(r, r) isa Array{Bool}
   end