layernorm: remove redundant arguments

Author: OneAdder

src/nf/nf_layernorm.f90

@@ -37,21 +37,17 @@ module nf_layernorm_layer
   end type layernorm_layer
 
   interface layernorm_layer
-    module function layernorm_layer_cons(sequence_length, model_dimension) &
+    module function layernorm_layer_cons() &
       result(res)
-      integer, intent(in) :: sequence_length, model_dimension
       type(layernorm_layer) :: res
     end function layernorm_layer_cons
   end interface layernorm_layer
 
 contains
-  module function layernorm_layer_cons(sequence_length, model_dimension) &
+  module function layernorm_layer_cons() &
     result(res)
-    integer, intent(in) :: sequence_length, model_dimension
     type(layernorm_layer) :: res
 
-    res % sequence_length = sequence_length
-    res % model_dimension = model_dimension
     res % eps = 1e-5
   end function layernorm_layer_cons
 
@@ -141,6 +137,12 @@ module subroutine init(self, input_shape)
     class(layernorm_layer), intent(in out) :: self
     integer, intent(in) :: input_shape(:)
 
+    if (size(input_shape) /= 2) then
+      error stop "LayerNorm Layer accepts 2D input"
+    end if
+    self % sequence_length = input_shape(1)
+    self % model_dimension = input_shape(2)
+
     ! default initialization from PyTorch
     allocate(self % gamma(self % model_dimension))
     self % gamma = 1.

test/test_layernorm.f90

@@ -8,12 +8,19 @@ program test_layernorm
   real :: sample_input(3, 4) = reshape([0.0, 10.1, 0.2, 10.3, 0.4, 10.5, 0.6, 10.7, 10.8, 0.9, 0.11, 0.12], [3, 4])
   real :: sample_gradient(3, 4) = reshape([0.1, 3., 2., 0.1, 3., 3., 0.1, 2., 0.1, 3., 0.1, 3.], [3, 4])
 
-  layernorm = layernorm_layer(3, 4)
-  call layernorm % init([0])
+  layernorm = layernorm_layer()
+  call layernorm % init([3, 4])
 
   call test_layernorm_forward(layernorm, sample_input, ok)
   call test_layernorm_backward(layernorm, sample_input, sample_gradient, ok)
 
+  if (ok) then
+    print '(a)', 'test_layernorm_layer: All tests passed.'
+  else
+    write(stderr, '(a)') 'test_layernorm_layer: One or more tests failed.'
+    stop 1
+  end if
+
 contains
   subroutine test_layernorm_forward(layernorm, input, ok)
     type(layernorm_layer), intent(in out) :: layernorm