Making indexing match regular arrays more closely

Author: TLipede

Project.toml

@@ -9,6 +9,7 @@ ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 
 [compat]
 Adapt = "3.0"
@@ -18,6 +19,7 @@ MLUtils = "0.2"
 NNlib = "0.8"
 
 [extras]
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]

src/array.jl

@@ -34,15 +34,18 @@ _isonehot(x::Base.ReshapedArray{<:Any, <:Any, <:OneHotArray{<:Any, L}}) where L
 
 Base.size(x::OneHotArray{<:Any, L}) where L = (Int(L), size(x.indices)...)
 
-_onehotindex(x, i) = (x == i)
+function Base.getindex(x::OneHotArray{<:Any, <:Any, N}, i::Integer, I::Vararg{Any, N}) where N
+  @boundscheck checkbounds(x, i, I...)
+  return x.indices[I...] .== i
+end
 
-Base.getindex(x::OneHotVector, i::Integer) = _onehotindex(x.indices, i)
-Base.getindex(x::OneHotVector{T, L}, ::Colon) where {T, L} = x
+function Base.getindex(x::OneHotArray{<:Any, L}, ::Colon, I...) where L
+  @boundscheck checkbounds(x, :, I...)
+  return OneHotArray(x.indices[I...], L)
+end
 
-Base.getindex(x::OneHotArray, i::Integer, I...) = _onehotindex.(x.indices[I...], i)
-Base.getindex(x::OneHotArray{<:Any, L}, ::Colon, I...) where L = OneHotArray(x.indices[I...], L)
-Base.getindex(x::OneHotArray{<:Any, <:Any, <:Any, N}, ::Vararg{Colon, N}) where N = x
-Base.getindex(x::OneHotArray, I::CartesianIndex{N}) where N = x[I[1], Tuple(I)[2:N]...]
+Base.getindex(x::OneHotArray, ::Colon) = BitVector(reshape(x, :))
+Base.getindex(x::OneHotArray{<:Any, <:Any, N}, ::Colon, ::Vararg{Colon, N}) where N = x
 
 function Base.showarg(io::IO, x::OneHotArray, toplevel)
   print(io, ndims(x) == 1 ? "OneHotVector(" : ndims(x) == 2 ? "OneHotMatrix(" : "OneHotArray(")

test/array.jl

@@ -1,3 +1,7 @@
+using Random
+
+Random.seed!(0)
+
 ov = OneHotVector(rand(1:10), 10)
 ov2 = OneHotVector(rand(1:11), 11)
 om = OneHotMatrix(rand(1:10, 5), 10)
@@ -21,6 +25,7 @@ end
   @test om[:, 3] == OneHotVector(om.indices[3], 10)
   @test om[3, :] == (om.indices .== 3)
   @test om[:, :] == om
+  @test om[:] == reshape(om, :)
 
   # array indexing
   @test oa[3, 3, 3] == (oa.indices[3, 3] == 3)
@@ -29,9 +34,22 @@ end
   @test oa[3, :, :] == (oa.indices .== 3)
   @test oa[:, 3, :] == OneHotMatrix(oa.indices[3, :], 10)
   @test oa[:, :, :] == oa
+  @test oa[:] == reshape(oa, :)
 
   # cartesian indexing
   @test oa[CartesianIndex(3, 3, 3)] == oa[3, 3, 3]
+
+  # linear indexing
+  @test om[9] == true
+  @test om[10] == false
+  @test om[11] == om[1, 2]
+  @test oa[52] == oa[2, 1, 2]
+
+  # bounds checks
+  @test_throws BoundsError ov[0]
+  @test_throws BoundsError om[2, -1]
+  @test_throws BoundsError oa[11, 5, 5]
+  @test_throws BoundsError oa[:, :]
 end
 
 @testset "Concatenating" begin