feat: Implementing reset state for RNN

Author: castelao

src/nf/nf_layer.f90

@@ -32,6 +32,7 @@ module nf_layer
     procedure :: set_params
     procedure :: init
     procedure :: print_info
+    procedure :: reset
 
     ! Specific subroutines for different array ranks
     procedure, private :: backward_1d
@@ -153,6 +154,10 @@ module subroutine set_params(self, params)
         !! Parameters of this layer
     end subroutine set_params
 
+    module subroutine reset(self)
+      class(layer), intent(in out) :: self
+    end subroutine reset
+
   end interface
 
 end module nf_layer

src/nf/nf_layer_submodule.f90

@@ -442,4 +442,14 @@ module subroutine set_params(self, params)
 
   end subroutine set_params
 
+  module subroutine reset(self)
+    class(layer), intent(in out) :: self
+
+    select type (this_layer => self % p)
+      type is (rnn_layer)
+        call this_layer % reset()
+    end select
+
+  end subroutine reset
+
 end submodule nf_layer_submodule

src/nf/nf_network.f90

@@ -26,6 +26,7 @@ module nf_network
     procedure :: get_params
     procedure :: print_info
     procedure :: set_params
+    procedure :: reset
     procedure :: train
     procedure :: update
 
@@ -242,6 +243,13 @@ module subroutine update(self, optimizer, batch_size)
         !! Set to `size(input_data, dim=2)` for a batch gradient descent.
     end subroutine update
 
+    module subroutine reset(self)
+      !! Reset network state
+      !!
+      !! Currently only affect RNN layer type
+      class(network), intent(in out) :: self
+    end subroutine reset
+
   end interface
 
 end module nf_network

src/nf/nf_network_submodule.f90

@@ -734,4 +734,18 @@ module subroutine update(self, optimizer, batch_size)
 
   end subroutine update
 
+  module subroutine reset(self)
+    class(network), intent(in out) :: self
+    integer :: n, num_layers
+
+    num_layers = size(self % layers)
+    do n = 2, num_layers
+      select type(this_layer => self % layers(n) % p)
+        type is(rnn_layer)
+          call self % layers(n) % reset()
+      end select
+    end do
+
+  end subroutine reset
+
 end submodule nf_network_submodule

src/nf/nf_rnn_layer.f90

@@ -40,6 +40,7 @@ module nf_rnn_layer
     procedure :: get_params
     procedure :: init
     procedure :: set_params
+    procedure :: reset
 
   end type rnn_layer
 
@@ -126,6 +127,14 @@ module subroutine init(self, input_shape)
         !! Shape of the input layer
     end subroutine init
 
+    module subroutine reset(self)
+      !! Reset layer state
+      !!
+      !! Currently reset state to zero but might be worth reconsidering it
+      !! in the future.
+      class(rnn_layer), intent(in out) :: self
+    end subroutine reset
+
   end interface
 
 end module nf_rnn_layer

src/nf/nf_rnn_layer_submodule.f90

@@ -166,4 +166,11 @@ module subroutine init(self, input_shape)
 
   end subroutine init
 
+  module subroutine reset(self)
+    class(rnn_layer), intent(in out) :: self
+
+    self % state = 0
+
+  end subroutine reset
+
 end submodule nf_rnn_layer_submodule