Assign layer_shape for maxpool2d in layer % init()

Author: milancurcic

src/nf_layer_submodule.f90

@@ -4,6 +4,7 @@
   use nf_dense_layer, only: dense_layer
   use nf_input1d_layer, only: input1d_layer
   use nf_input3d_layer, only: input3d_layer
+  use nf_maxpool2d_layer, only: maxpool2d_layer
 
   implicit none
 
@@ -109,9 +110,13 @@ impure elemental module subroutine init(self, input)
       call this_layer % init(input % layer_shape)
     end select
 
-    ! The shape of a conv2d layer is not known until we receive an input layer.
-    select type(this_layer => self % p); type is(conv2d_layer)
-      self % layer_shape = shape(this_layer % output)
+    ! The shape of conv2d or maxpool2d layers is not known
+    ! until we receive an input layer.
+    select type(this_layer => self % p)
+      type is(conv2d_layer)
+        self % layer_shape = shape(this_layer % output)
+      type is(maxpool2d_layer)
+        self % layer_shape = shape(this_layer % output)
     end select
 
     self % input_layer_shape = input % layer_shape 

test/test_maxpool2d_layer.f90

@@ -23,6 +23,24 @@ program test_maxpool2d_layer
     write(stderr, '(a)') 'maxpool2d layer should not be marked as initialized yet.. failed'
   end if
 
+  input_layer = input([3, 32, 32])
+  call maxpool_layer % init(input_layer)
+
+  if (.not. maxpool_layer % initialized) then
+    ok = .false.
+    write(stderr, '(a)') 'maxpool2d layer should now be marked as initialized.. failed'
+  end if
+
+  if (.not. all(maxpool_layer % input_layer_shape == [3, 32, 32])) then
+    ok = .false.
+    write(stderr, '(a)') 'maxpool2d layer input layer shape should be correct.. failed'
+  end if
+
+  if (.not. all(maxpool_layer % layer_shape == [3, 16, 16])) then
+    ok = .false.
+    write(stderr, '(a)') 'maxpool2d layer input layer shape should be correct.. failed'
+  end if
+
   if (ok) then
     print '(a)', 'test_maxpool2d_layer: All tests passed.'
   else