Fast path onehotbatch(::Vector{Int}, ::UnitRange) (#27)

* add a fast path

* add an error check

* fixup, add tests

* fix 1.6

Project.toml

@@ -1,6 +1,6 @@
 name = "OneHotArrays"
 uuid = "0b1bfda6-eb8a-41d2-88d8-f5af5cad476f"
-version = "0.2.1"
+version = "0.2.2"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/onehot.jl

@@ -100,6 +100,16 @@ function _onehotbatch(data, labels, default)
   return OneHotArray(indices, length(labels))
 end
 
+function onehotbatch(data::AbstractArray{<:Integer}, labels::AbstractUnitRange{<:Integer})
+  # lo, hi = extrema(data)  # fails on Julia 1.6
+  lo, hi = minimum(data), maximum(data)
+  lo < first(labels) && error("Value $lo not found in labels")
+  hi > last(labels) && error("Value $hi not found in labels")
+  offset = 1 - first(labels)
+  indices = UInt32.(data .+ offset)
+  return OneHotArray(indices, length(labels))
+end
+
 """
     onecold(y::AbstractArray, labels = 1:size(y,1))
 

test/gpu.jl

@@ -26,6 +26,16 @@ end
   @test_broken gradient(A -> sum(A * y), gA)[1] isa CuArray  # fails with JLArray, bug in Zygote?
 end
 
+@testset "onehotbatch(::CuArray, ::UnitRange)" begin
+  y1 = onehotbatch([1, 3, 0, 2], 0:9) |> cu
+  y2 = onehotbatch([1, 3, 0, 2] |> cu, 0:9)
+  @test y1.indices == y2.indices
+  @test_broken y1 == y2
+
+  @test_throws Exception onehotbatch([1, 3, 0, 2] |> cu, 1:10)
+  @test_throws Exception onehotbatch([1, 3, 0, 2] |> cu, -2:2)
+end
+
 @testset "onecold gpu" begin
   y = onehotbatch(ones(3), 1:10) |> cu;
   l = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j']

test/onehot.jl

@@ -27,6 +27,12 @@
   @test onecold(onehot(-0.0, floats)) == 2  # as it uses isequal
   @test onecold(onehot(Inf, floats)) == 5
 
+  # UnitRange fast path
+  @test onehotbatch([1,3,0,4], 0:4) == onehotbatch([1,3,0,4], Tuple(0:4))
+  @test onehotbatch([2 3 7 4], 2:7) == onehotbatch([2 3 7 4], Tuple(2:7))
+  @test_throws Exception onehotbatch([2, -1], 0:4)
+  @test_throws Exception onehotbatch([2, 5], 0:4)
+
   # inferrabiltiy tests
   @test @inferred(onehot(20, 10:10:30)) == [false, true, false]
   @test @inferred(onehot(40, (10,20,30), 20)) == [false, true, false]