Merge pull request #7 from TLipede/indexing-changes

Author: darsnack

Project.toml

@@ -12,8 +12,8 @@ NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 
 [compat]
 Adapt = "3.0"
-ChainRulesCore = "1.13"
 CUDA = "3.8"
+ChainRulesCore = "1.13"
 MLUtils = "0.2"
 NNlib = "0.8"
 

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

@@ -21,6 +21,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 +30,20 @@ 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[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