embedding_layer: implementation of embedding layer

src/nf/nf_embedding_layer.f90

@@ -9,6 +9,11 @@ module nf_embedding_layer
   public :: embedding_layer
 
   type, extends(base_layer) :: embedding_layer
+    !! Embedding Layer
+    !! Stores inputs as a trainable lookup table. Inputs are
+    !! integer indicies in a dictionary of `vocab_size`.
+    !! This layer converts them into a table of shape
+    !! (`sequence_length`, `model_dimension`)
     integer :: sequence_length, vocab_size, model_dimension
 
     real, allocatable :: weights(:, :)
@@ -29,24 +34,25 @@ module nf_embedding_layer
   end type embedding_layer
 
   interface embedding_layer
-    module function embedding_layer_cons(&
-        sequence_length, vocab_size, model_dimension&
-    ) result(res)
-      integer, intent(in) :: sequence_length, vocab_size, model_dimension
+    module function embedding_layer_cons(vocab_size, model_dimension) result(res)
+      integer, intent(in) :: vocab_size, model_dimension
       type(embedding_layer) :: res
     end function embedding_layer_cons
   end interface embedding_layer
 
   interface
     pure module subroutine forward(self, input)
+      !! Get vectors by indicis in the dictionary
       class(embedding_layer), intent(in out) :: self
       integer, intent(in) :: input(:)
     end subroutine forward
 
     pure module subroutine backward(self, input, gradient)
+      !! Update gradient at `input` indices
+      !! dw_i = W_i + d_output_i
       class(embedding_layer), intent(in out) :: self
       integer, intent(in) :: input(:)
-      real, intent(in) :: gradient(:)
+      real, intent(in) :: gradient(:, :)
     end subroutine backward
 
     module subroutine init(self, input_shape)

src/nf/nf_embedding_submodule.f90

@@ -2,21 +2,20 @@
   use nf_base_layer, only: base_layer
   implicit none
 contains
-  module function embedding_layer_cons(&
-      sequence_length, vocab_size, model_dimension&
-  ) result(res)
-    integer, intent(in) :: sequence_length, vocab_size, model_dimension
+  module function embedding_layer_cons(vocab_size, model_dimension) result(res)
+    integer, intent(in) :: vocab_size, model_dimension
     type(embedding_layer) :: res
 
     res % vocab_size = vocab_size
     res % model_dimension = model_dimension
-    res % sequence_length = sequence_length
   end function embedding_layer_cons
 
   module subroutine init(self, input_shape)
     class(embedding_layer), intent(in out) :: self
     integer, intent(in) :: input_shape(:)
 
+    self % sequence_length = input_shape(1)
+
     allocate(self % output(self % sequence_length, self % model_dimension))
     allocate(self % gradient(self % sequence_length, self % vocab_size))
 
@@ -30,32 +29,34 @@ end subroutine init
   pure module subroutine forward(self, input)
     class(embedding_layer), intent(in out) :: self
     integer, intent(in) :: input(:)
-    integer :: i
+    integer :: i, index
 
     do concurrent(i = 1: self % sequence_length)
-      self % output(i, :) = self % weights(input(i), :)
+      index = input(i)
+      if (index > size(self % weights, 1)) then
+        index = 1
+      end if
+      self % output(i, :) = self % weights(index, :)
     end do
   end subroutine forward
 
   pure module subroutine backward(self, input, gradient)
     class(embedding_layer), intent(in out) :: self
     integer, intent(in) :: input(:)
-    real, intent(in) :: gradient(:)
-    real :: db(self % model_dimension)
-    real :: dw(self % vocab_size, self % model_dimension)
+    real, intent(in) :: gradient(:, :)
     integer :: i
+
+    do concurrent(i = 1: self % sequence_length)
+      self % dw(input(i), :) = self % dw(input(i), :) + gradient(i, :)
+    end do
   end subroutine backward
 
   pure module function get_num_params(self) result(num_params)
     class(embedding_layer), intent(in) :: self
     integer :: num_params
-
-    ! Number of weigths times number of biases
-    num_params = self % vocab_size * self % model_dimension + self % model_dimension
-
+    num_params = self % vocab_size * self % model_dimension
   end function get_num_params
 
-
   module function get_params(self) result(params)
     class(embedding_layer), intent(in), target :: self
     real, allocatable :: params(:)
@@ -65,7 +66,6 @@ module function get_params(self) result(params)
     params = [w_]
   end function get_params
 
-
   module function get_gradients(self) result(gradients)
     class(embedding_layer), intent(in), target :: self
     real, allocatable :: gradients(:)
@@ -75,7 +75,6 @@ module function get_gradients(self) result(gradients)
     gradients = [dw_]
   end function get_gradients
 
-
   module subroutine set_params(self, params)
     class(embedding_layer), intent(in out) :: self
     real, intent(in), target :: params(:)

test/test_embedding_layer.f90

@@ -5,10 +5,36 @@ program test_embedding_layer
 
   logical :: ok = .true.
   integer :: sample_input(3) = [2, 1, 3]
+  real :: sample_gradient(3, 2) = reshape([0.1, 0.2, 0.3, 0.4, 0.6, 0.6], [3, 2])
+  real :: output_flat(6)
+  real :: expected_output_flat(6) = reshape([0.3, 0.1, 0.5, 0.4, 0.2, 0.6], [6])
+  real :: dw_flat(8)
+  real :: expected_dw_flat(8) = reshape([0.2, 0.1, 0.3, 0., 0.6, 0.4, 0.6, 0.], [8])
   type(embedding_layer) :: embedding
 
-  embedding = embedding_layer(sequence_length=3, vocab_size=4, model_dimension=2)
-  call embedding % init([0])
+  embedding = embedding_layer(vocab_size=4, model_dimension=2)
+  call embedding % init([3])
   embedding % weights = reshape([0.1, 0.3, 0.5, 0.7, 0.2, 0.4, 0.6, 0.8], [4, 2])
+
   call embedding % forward(sample_input)
-end program test_embedding_layer
+
+  output_flat = reshape(embedding % output, [6])
+  if (.not. all(output_flat.eq.expected_output_flat)) then
+    ok = .false.
+    write(stderr, '(a)') 'forward returned incorrect values.. failed'
+  end if
+
+  call embedding % backward(sample_input, sample_gradient)
+  dw_flat = reshape(embedding % dw, shape(dw_flat))
+  if (.not. all(dw_flat.eq.expected_dw_flat)) then
+    ok = .false.
+    write(stderr, '(a)') 'backward returned incorrect dw values.. failed'
+  end if
+
+  if (ok) then
+    print '(a)', 'test_embedding_layer: All tests passed.'
+  else
+    write(stderr, '(a)') 'test_embedding_layer: One or more tests failed.'
+    stop 1
+  end if
+end program test_embedding_layer