Merge pull request #89 from theogf/VecVecs

Added VecOfVecs and RowVecs

Author: theogf

src/utils.jl

@@ -9,11 +9,20 @@ macro check_args(K, param, cond, desc=string(cond))
     end
 end
 
+function vec_of_vecs(X::AbstractMatrix; obsdim::Int = 2)
+    @assert obsdim ∈ (1, 2) "obsdim should be 1 or 2"
+    if obsdim == 1
+        RowVecs(X)
+    else
+        ColVecs(X)
+    end
+end
 
 """
     ColVecs(X::AbstractMatrix)
 
-A lightweight box for an `AbstractMatrix` to make it behave like a vector of vectors.
+A lightweight wrapper for an `AbstractMatrix` to make it behave like a vector of vectors.
+Each vector represents a column of the matrix
 """
 struct ColVecs{T, TX<:AbstractMatrix{T}, S} <: AbstractVector{S}
     X::TX
@@ -27,6 +36,24 @@ Base.size(D::ColVecs) = (size(D.X, 2),)
 Base.getindex(D::ColVecs, i::Int) = view(D.X, :, i)
 Base.getindex(D::ColVecs, i) = ColVecs(view(D.X, :, i))
 
+"""
+    RowVecs(X::AbstractMatrix)
+
+A lightweight wrapper for an `AbstractMatrix` to make it behave like a vector of vectors.
+Each vector represents a row of the matrix
+"""
+struct RowVecs{T, TX<:AbstractMatrix{T}, S} <: AbstractVector{S}
+    X::TX
+    function RowVecs(X::TX) where {T, TX<:AbstractMatrix{T}}
+        S = typeof(view(X, 1, :))
+        new{T, TX, S}(X)
+    end
+end
+
+Base.size(D::RowVecs) = (size(D.X, 1),)
+Base.getindex(D::RowVecs, i::Int) = view(D.X, i, :)
+Base.getindex(D::RowVecs, i) = RowVecs(view(D.X, i, :))
+
 # Take highest Float among possibilities
 # function promote_float(Tₖ::DataType...)
 #     if length(Tₖ) == 0

src/zygote_adjoints.jl

@@ -13,6 +13,15 @@ end
     return ColVecs(X), back
 end
 
+@adjoint function RowVecs(X::AbstractMatrix)
+    back(Δ::NamedTuple) = (Δ.X,)
+    back(Δ::AbstractMatrix) = (Δ,)
+    function back(Δ::AbstractVector{<:AbstractVector{<:Real}})
+        throw(error("In slow method"))
+    end
+    return RowVecs(X), back
+end
+
 # @adjoint function evaluate(s::Sinus, x::AbstractVector, y::AbstractVector)
 #   d = evaluate(s, x, y)
 #   s = sum(sin.(π*(x-y)))

test/runtests.jl

@@ -1,7 +1,7 @@
+using KernelFunctions
 using Distances
 using FiniteDifferences
 using Flux
-using KernelFunctions
 using Kronecker
 using LinearAlgebra
 using PDMats

test/utils.jl

@@ -1,8 +1,11 @@
 @testset "utils" begin
-    using KernelFunctions: ColVecs
-    rng, N, D = MersenneTwister(123456), 10, 2
+    using KernelFunctions: vec_of_vecs, ColVecs, RowVecs
+    rng, N, D = MersenneTwister(123456), 10, 4
     x, X = randn(rng, N), randn(rng, D, N)
-
+    @testset "VecOfVecs" begin
+        @test vec_of_vecs(X, obsdim = 2) == ColVecs(X)
+        @test vec_of_vecs(X, obsdim = 1) == RowVecs(X)
+    end
     # Test Matrix data sets.
     @testset "ColVecs" begin
         DX = ColVecs(X)
@@ -22,6 +25,29 @@
             DX, back = Zygote.pullback(ColVecs, X)
             @test back((X=ones(size(X)),))[1] == ones(size(X))
 
+            @test Zygote.pullback(DX->DX.X, DX)[1] == X
+            X_, back = Zygote.pullback(DX->DX.X, DX)
+            @test back(ones(size(X)))[1].X == ones(size(X))
+        end
+    end
+    @testset "RowVecs" begin
+        DX = RowVecs(X)
+        @test DX == DX
+        @test size(DX) == (D,)
+        @test length(DX) == D
+        @test getindex(DX, 2) isa AbstractVector
+        @test getindex(DX, 2) == X[2, :]
+        @test getindex(DX, 1:3) isa RowVecs
+        @test getindex(DX, 1:3) == RowVecs(X[1:3, :])
+        @test getindex(DX, :) == RowVecs(X)
+        @test eachindex(DX) == 1:D
+        @test first(DX) == X[1, :]
+
+        let
+            @test Zygote.pullback(RowVecs, X)[1] == DX
+            DX, back = Zygote.pullback(RowVecs, X)
+            @test back((X=ones(size(X)),))[1] == ones(size(X))
+
             @test Zygote.pullback(DX->DX.X, DX)[1] == X
             X_, back = Zygote.pullback(DX->DX.X, DX)
             @test back(ones(size(X)))[1].X == ones(size(X))