Merge pull request #670 from ldeso/add-cartesian-embedding-methods

Add Cartesian Embedding methods

src/layers/basic.jl

@@ -451,7 +451,8 @@ end
     Embedding(in_dims => out_dims; init_weight=randn32)
 
 A lookup table that stores embeddings of dimension `out_dims` for a vocabulary of size
-`in_dims`.
+`in_dims`. When the vocabulary is multi-dimensional, the input is expected to be a tuple
+of Cartesian indices.
 
 This layer is often used to store word embeddings and retrieve them using indices.
 
@@ -461,19 +462,22 @@ This layer is often used to store word embeddings and retrieve them using indice
 
 ## Arguments
 
-  - `in_dims`: number of input dimensions
+  - `in_dims`: number(s) of input dimensions
   - `out_dims`: number of output dimensions
 
 ## Keyword Arguments
 
   - `init_weight`: initializer for the weight matrix
-    (`weight = init_weight(rng, out_dims, in_dims)`)
+    (`weight = init_weight(rng, out_dims, in_dims...)`)
 
 ## Input
 
   - Integer OR
   - Abstract Vector of Integers OR
-  - Abstract Array of Integers
+  - Abstract Array of Integers OR
+  - Tuple of Integers OR
+  - Tuple of Abstract Vectors of Integers OR
+  - Tuple of Abstract Arrays of Integers
 
 ## Returns
 
@@ -482,17 +486,19 @@ This layer is often used to store word embeddings and retrieve them using indice
   - Empty `NamedTuple()`
 """
 @concrete struct Embedding <: AbstractExplicitLayer
-    in_dims::Int
+    in_dims
     out_dims::Int
     init_weight
 end
 
-function Embedding((in_dims, out_dims)::Pair{<:Integer, <:Integer}; init_weight=randn32)
+function Embedding(
+        (in_dims, out_dims)::Pair{<:Union{Integer, NTuple{<:Any, <:Integer}}, <:Integer};
+        init_weight=randn32)
     return Embedding(in_dims, out_dims, init_weight)
 end
 
 function initialparameters(rng::AbstractRNG, e::Embedding)
-    return (weight=e.init_weight(rng, e.out_dims, e.in_dims),)
+    return (weight=e.init_weight(rng, e.out_dims, e.in_dims...),)
 end
 
 (e::Embedding)(x::Integer, ps, st::NamedTuple) = view(ps.weight, :, x), st
@@ -502,6 +508,22 @@ end
 function (e::Embedding)(x::AbstractArray{<:Integer}, ps, st::NamedTuple)
     return reshape(e(vec(x), ps, st)[1], :, size(x)...), st
 end
+function (e::Embedding)(x::NTuple{<:Any, <:Integer}, ps, st::NamedTuple)
+    view(ps.weight, :, x...), st
+end
+function (e::Embedding)(x::NTuple{<:Any, <:AbstractVector{<:Integer}}, ps, st::NamedTuple)
+    sizes = size.(x)
+    @argcheck allequal(sizes) DimensionMismatch("Input vectors must have the same shape")
+    return NNlib.gather(ps.weight, x...), st
+end
+function (e::Embedding)(x::NTuple{<:Any, <:AbstractArray{<:Integer}}, ps, st::NamedTuple)
+    sizes = size.(x)
+    @argcheck allequal(sizes) DimensionMismatch("Input arrays must have the same shape")
+    return reshape(e(vec.(x), ps, st)[1], :, first(sizes)...), st
+end
+function (e::Embedding)(x::Tuple{}, ps, st::NamedTuple)
+    throw(ArgumentError("Input tuple must contain at least one element"))
+end
 
 function Base.show(io::IO, e::Embedding)
     return print(io, "Embedding(", e.in_dims, " => ", e.out_dims, ")")

test/layers/basic_tests.jl

@@ -334,34 +334,77 @@ end
     rng = get_stable_rng(12345)
 
     @testset "$mode" for (mode, aType, device, ongpu) in MODES
-        vocab_size, embed_size = 10, 4
-        layer = Embedding(vocab_size => embed_size)
-        __display(layer)
-        ps, st = Lux.setup(rng, layer) .|> device
+        @testset "Linear indices" begin
+            vocab_size, embed_size = 10, 4
+            layer = Embedding(vocab_size => embed_size)
+            __display(layer)
+            ps, st = Lux.setup(rng, layer) .|> device
+
+            @test size(ps.weight) == (embed_size, vocab_size)
+
+            @test LuxCore.outputsize(layer) == (4,)
+
+            x = rand(1:vocab_size, 1)[1]
+            y, st_ = layer(x, ps, st)
+            @test size(layer(x, ps, st)[1]) == (embed_size,)
+            @test y == ps.weight[:, x]
+
+            @jet layer(x, ps, st)
+
+            x = rand(1:vocab_size, 3) |> aType
+            y, st_ = layer(x, ps, st)
+            @test y isa aType{Float32}
+            @test y == ps.weight[:, x]
+
+            @jet layer(x, ps, st)
+
+            x = rand(1:vocab_size, 3, 4) |> aType
+            y, st_ = layer(x, ps, st)
+            @test y isa aType{Float32, 3}
+            @test size(y) == (embed_size, 3, 4)
+
+            @jet layer(x, ps, st)
+        end
+
+        @testset "Cartesian indices" begin
+            vocab_size, embed_size = (5, 2), 4
+            layer = Embedding(vocab_size => embed_size)
+            __display(layer)
+            ps, st = Lux.setup(rng, layer) .|> device
+
+            @test size(ps.weight) == (embed_size, vocab_size...)
 
-        @test size(ps.weight) == (embed_size, vocab_size)
+            @test LuxCore.outputsize(layer) == (4,)
 
-        @test LuxCore.outputsize(layer) == (4,)
+            x = (rand(1:vocab_size[1], 1)[1], rand(1:vocab_size[2], 1)[1])
+            y, st_ = layer(x, ps, st)
+            @test size(layer(x, ps, st)[1]) == (embed_size,)
+            @test y == ps.weight[:, x...]
 
-        x = rand(1:vocab_size, 1)[1]
-        y, st_ = layer(x, ps, st)
-        @test size(layer(x, ps, st)[1]) == (embed_size,)
-        @test y == ps.weight[:, x]
+            @jet layer(x, ps, st)
+
+            x = (rand(1:vocab_size[1], 3), rand(1:vocab_size[2], 3)) .|> aType
+            y, st_ = layer(x, ps, st)
+            @test y isa aType{Float32}
+            @test y == ps.weight[:, CartesianIndex.(x...)]
+
+            @jet layer(x, ps, st)
 
-        @jet layer(x, ps, st)
+            x = (rand(1:vocab_size[1], 3, 4), rand(1:vocab_size[2], 3, 4)) .|> aType
+            y, st_ = layer(x, ps, st)
+            @test y isa aType{Float32, 3}
+            @test size(y) == (embed_size, 3, 4)
 
-        x = rand(1:vocab_size, 3) |> aType
-        y, st_ = layer(x, ps, st)
-        @test y isa aType{Float32}
-        @test y == ps.weight[:, x]
+            @jet layer(x, ps, st)
 
-        @jet layer(x, ps, st)
+            x = (rand(1:vocab_size[1], 3), rand(1:vocab_size[2], 4)) .|> aType
+            @test_throws DimensionMismatch layer(x, ps, st)
 
-        x = rand(1:vocab_size, 3, 4) |> aType
-        y, st_ = layer(x, ps, st)
-        @test y isa aType{Float32, 3}
-        @test size(y) == (embed_size, 3, 4)
+            x = (rand(1:vocab_size[1], 3, 4), rand(1:vocab_size[2], 4, 5)) .|> aType
+            @test_throws DimensionMismatch layer(x, ps, st)
 
-        @jet layer(x, ps, st)
+            x = ()
+            @test_throws ArgumentError layer(x, ps, st)
+        end
     end
 end