Relaxed condition getindex(X, i)

Author: theogf

src/utils.jl

@@ -11,7 +11,7 @@ end
 
 
 """
-    ColVecs{T, TX<:AbstractMatrix}
+    ColVecs(X::AbstractMatrix)
 
 A lightweight box for an `AbstractMatrix` to make it behave like a vector of vectors.
 """
@@ -25,7 +25,7 @@ end
 
 Base.size(D::ColVecs) = (size(D.X, 2),)
 Base.getindex(D::ColVecs, i::Int) = view(D.X, :, i)
-Base.getindex(D::ColVecs, i::AbstractVector{Int}) = ColVecs(view(D.X, :, i))
+Base.getindex(D::ColVecs, i) = ColVecs(view(D.X, :, i))
 
 # Take highest Float among possibilities
 # function promote_float(Tₖ::DataType...)

test/utils.jl

@@ -13,10 +13,10 @@
         @test getindex(DX, 5) == X[:, 5]
         @test getindex(DX, 1:2:6) isa ColVecs
         @test getindex(DX, 1:2:6) == ColVecs(X[:, 1:2:6])
+        @test getindex(DX, :) == ColVecs(X)
         @test eachindex(DX) == 1:N
         @test first(DX) == X[:, 1]
 
-
         let
             @test Zygote.pullback(ColVecs, X)[1] == DX
             DX, back = Zygote.pullback(ColVecs, X)