embedding_layer: positional encoding

Author: OneAdder

src/nf/nf_embedding_layer.f90

@@ -15,6 +15,7 @@ module nf_embedding_layer
     !! This layer converts them into a table of shape
     !! (`sequence_length`, `model_dimension`)
     integer :: sequence_length, vocab_size, model_dimension
+    logical :: positional
 
     real, allocatable :: weights(:, :)
     real, allocatable :: output(:, :)
@@ -24,6 +25,7 @@ module nf_embedding_layer
 
     procedure :: backward
     procedure :: forward
+    procedure :: positional_encoding
     procedure :: init
     procedure :: get_num_params
     procedure :: get_params
@@ -33,8 +35,9 @@ module nf_embedding_layer
   end type embedding_layer
 
   interface embedding_layer
-    module function embedding_layer_cons(vocab_size, model_dimension) result(res)
+    module function embedding_layer_cons(vocab_size, model_dimension, positional) result(res)
       integer, intent(in) :: vocab_size, model_dimension
+      logical, optional :: positional
       type(embedding_layer) :: res
     end function embedding_layer_cons
   end interface embedding_layer
@@ -54,6 +57,11 @@ pure module subroutine backward(self, input, gradient)
       real, intent(in) :: gradient(:, :)
     end subroutine backward
 
+    pure module subroutine positional_encoding(self, pos)
+      class(embedding_layer), intent(in out) :: self
+      integer, intent(in) :: pos
+    end subroutine positional_encoding
+
     module subroutine init(self, input_shape)
       class(embedding_layer), intent(in out) :: self
       integer, intent(in) :: input_shape(:)

src/nf/nf_embedding_layer_submodule.f90

@@ -2,12 +2,18 @@
   use nf_base_layer, only: base_layer
   implicit none
 contains
-  module function embedding_layer_cons(vocab_size, model_dimension) result(res)
+  module function embedding_layer_cons(vocab_size, model_dimension, positional) result(res)
     integer, intent(in) :: vocab_size, model_dimension
+    logical, optional :: positional
     type(embedding_layer) :: res
 
     res % vocab_size = vocab_size
     res % model_dimension = model_dimension
+    if (.not. present(positional)) then
+      res % positional = .false.
+    else
+      res % positional = positional
+    end if
   end function embedding_layer_cons
 
   module subroutine init(self, input_shape)
@@ -37,7 +43,12 @@ pure module subroutine forward(self, input)
       elseif (index == 0) then
         index = 1
       end if
+
       self % output(i, :) = self % weights(index, :)
+
+      if (self % positional) then
+        call self % positional_encoding(i)
+      end if
     end do
   end subroutine forward
 
@@ -52,6 +63,19 @@ pure module subroutine backward(self, input, gradient)
     end do
   end subroutine backward
 
+  pure module subroutine positional_encoding(self, pos)
+    class(embedding_layer), intent(in out) :: self
+    integer, intent(in) :: pos
+    integer :: i
+    real :: theta
+
+    do concurrent(i = 1: floor(real(self % model_dimension) / 2))
+      theta = (pos - 1) / 10000 ** (real(2 * (i-1)) / self % model_dimension)
+      self % output(pos, 2 * i - 1) = self % output(pos, 2 * i - 1) + sin(theta)
+      self % output(pos, 2 * i) = self % output(pos, 2 * i) + cos(theta)
+    end do
+  end subroutine positional_encoding
+
   pure module function get_num_params(self) result(num_params)
     class(embedding_layer), intent(in) :: self
     integer :: num_params