Add some checks to onecold (#25)

* add sanity checks

* Apply suggestions from code review

Co-authored-by: Brian Chen <ToucheSir@users.noreply.github.com>

Co-authored-by: Brian Chen <ToucheSir@users.noreply.github.com>

Author: mcabbott

Project.toml

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

src/onehot.jl

@@ -126,8 +126,18 @@ julia> onecold([ 1  0  0  1  0  1  0  1  0  0  1
 "abeacadabea"
 ```
 """
-onecold(y::AbstractVector, labels = 1:length(y)) = labels[argmax(y)]
+function onecold(y::AbstractVector, labels = 1:length(y))
+  nl = length(labels)
+  ny = length(y)
+  nl == ny || throw(DimensionMismatch("onecold got $nl labels for a vector of length $ny, these must agree"))
+  ymax, i = findmax(y)
+  ymax isa Number && isnan(ymax) && throw(ArgumentError("maximum value found by onecold is $ymax"))
+  labels[i]
+end
 function onecold(y::AbstractArray, labels = 1:size(y, 1))
+  nl = length(labels)
+  ny = size(y, 1)
+  nl == ny || throw(DimensionMismatch("onecold got $nl labels for an array with size(y, 1) == $ny, these must agree"))
   indices = _fast_argmax(y)
   xs = isbits(labels) ? indices : collect(indices) # non-bit type cannot be handled by CUDA
 

test/onehot.jl

@@ -52,6 +52,10 @@ end
   cold = onecold(hot, labels)
 
   @test cold == data
+
+  @test_throws DimensionMismatch onecold([0.3, 0.2], (:a, :b, :c))
+  @test_throws DimensionMismatch onecold([0.3, 0.2, 0.5, 0.99], (:a, :b, :c))
+  @test_throws ArgumentError onecold([0.3, NaN, 0.5], (:a, :b, :c))
 end
 
 @testset "onehotbatch indexing" begin