little modifies

Author: Riccardo231

src/nf.f90

@@ -3,7 +3,7 @@ module nf
   use nf_datasets_mnist, only: label_digits, load_mnist
   use nf_layer, only: layer
   use nf_layer_constructors, only: &
-    conv2d, dense, flatten, input, maxpool2d, reshape, locally_connected_1d
+    conv2d, dense, flatten, input, maxpool2d, reshape, reshape_generalized, locally_connected_1d
   use nf_loss, only: mse, quadratic
   use nf_metrics, only: corr, maxabs
   use nf_network, only: network

src/nf/nf_layer_constructors_submodule.f90

@@ -162,14 +162,16 @@ module function reshape(output_shape) result(res)
   end function reshape
 
   module function reshape_generalized(output_shape) result(res)
-    integer, intent(in) :: output_shape(:)
+    integer, intent(in) :: output_shape(:)  !! Always treat as an array
     type(layer) :: res
-
+    
     res % name = 'reshape_generalized'
     res % layer_shape = output_shape
 
     allocate(res % p, source=reshape_generalized_layer(output_shape))
-  
+
   end function reshape_generalized
+  
+  
 
 end submodule nf_layer_constructors_submodule

src/nf/nf_reshape_generalized.f90

@@ -1,77 +1,76 @@
 module nf_reshape_layer_generalized
 
-    !! This module provides the concrete reshape layer type.
-    !! It is used internally by the layer type.
-    !! It is not intended to be used directly by the user.
-  
-    use nf_base_layer, only: base_layer
-  
-    implicit none
-  
-    private
-    public :: reshape_generalized_layer
-  
-    type, extends(base_layer) :: reshape_generalized_layer
-  
-      !! Concrete implementation of a reshape layer type
-      !! It implements reshaping for arbitrary ranks.
-  
-      integer, allocatable :: input_shape(:)
-      integer, allocatable :: output_shape(:)
-      real, allocatable :: gradient(:)
-      real, allocatable :: output(:)
-  
-    contains
-  
-      procedure :: backward
-      procedure :: forward
-      procedure :: init
-  
-    end type reshape_generalized_layer
-  
-    interface reshape_generalized_layer
-      pure module function reshape_layer_cons(output_shape) result(res)
-        !! This function returns the `reshape_layer` instance.
-        integer, intent(in) :: output_shape(:)
-          !! The shape of the output
-        type(reshape_generalized_layer) :: res
-          !! reshape_layer instance
-      end function reshape_layer_cons
-    end interface reshape_generalized_layer
-  
-    interface
-  
-      pure module subroutine backward(self, input, gradient)
-        !! Apply the backward pass for the reshape layer.
-        !! This is just flattening to a rank-1 array.
-        class(reshape_generalized_layer), intent(in out) :: self
-          !! Dense layer instance
-        real, intent(in) :: input(:)
-          !! Input from the previous layer
-        real, intent(in) :: gradient(..)
-          !! Gradient from the next layer
-      end subroutine backward
-  
-      pure module subroutine forward(self, input)
-        !! Apply the forward pass for the reshape layer.
-        !! This is reshaping from input rank to output rank.
-        class(reshape_generalized_layer), intent(in out) :: self
-          !! Dense layer instance
-        real, intent(in) :: input(:)
-          !! Input from the previous layer
-      end subroutine forward
-  
-      module subroutine init(self, input_shape)
-        !! Initialize the layer data structures.
-        !!
-        !! This is a deferred procedure from the `base_layer` abstract type.
-        class(reshape_generalized_layer), intent(in out) :: self
-          !! Dense layer instance
-        integer, intent(in) :: input_shape(:)
-          !! Shape of the input layer
-      end subroutine init
-  
-    end interface
-  
-  end module nf_reshape_layer_generalized
-  
+  !! This module provides the concrete reshape layer type.
+  !! It is used internally by the layer type.
+  !! It is not intended to be used directly by the user.
+
+  use nf_base_layer, only: base_layer
+
+  implicit none
+
+  private
+  public :: reshape_generalized_layer
+
+  type, extends(base_layer) :: reshape_generalized_layer
+
+    !! Concrete implementation of a reshape layer type
+    !! It implements reshaping for arbitrary ranks.
+
+    integer, allocatable :: input_shape(:)
+    integer, allocatable :: output_shape(:)
+    real, allocatable :: gradient(:)
+    real, allocatable :: output(:)
+
+  contains
+
+    procedure :: backward
+    procedure :: forward
+    procedure :: init
+
+  end type reshape_generalized_layer
+
+  interface reshape_generalized_layer
+    pure module function reshape_layer_cons(output_shape) result(res)
+      !! This function returns the `reshape_layer` instance.
+      integer, intent(in) :: output_shape(:)
+        !! The shape of the output
+      type(reshape_generalized_layer) :: res
+        !! reshape_layer instance
+    end function reshape_layer_cons
+  end interface reshape_generalized_layer
+
+  interface
+
+    pure module subroutine backward(self, input, gradient)
+      !! Apply the backward pass for the reshape layer.
+      !! This is just flattening to a rank-1 array.
+      class(reshape_generalized_layer), intent(in out) :: self
+        !! Dense layer instance
+      real, intent(in) :: input(:)
+        !! Input from the previous layer
+      real, intent(in) :: gradient(..)
+        !! Gradient from the next layer
+    end subroutine backward
+
+    pure module subroutine forward(self, input)
+      !! Apply the forward pass for the reshape layer.
+      !! This is reshaping from input rank to output rank.
+      class(reshape_generalized_layer), intent(in out) :: self
+        !! Dense layer instance
+      real, intent(in) :: input(:)
+        !! Input from the previous layer
+    end subroutine forward
+
+    module subroutine init(self, input_shape)
+      !! Initialize the layer data structures.
+      !! This is a deferred procedure from the `base_layer` abstract type.
+      class(reshape_generalized_layer), intent(in out) :: self
+        !! Dense layer instance
+      integer, intent(in) :: input_shape(:)
+        !! Shape of the input layer
+    end subroutine init
+
+  end interface
+
+end module nf_reshape_layer_generalized
+

src/nf/nf_reshape_generalized_submodule.f90

@@ -9,10 +9,18 @@
   pure module function reshape_layer_cons(output_shape) result(res)
     integer, intent(in) :: output_shape(:)
     type(reshape_generalized_layer) :: res
-    allocate(res % output_shape(size(output_shape)))
-    res % output_shape = output_shape
+    
+    ! Check if output_shape is scalar (size 1)
+    if (size(output_shape) == 1) then
+        allocate(res % output_shape(1))
+        res % output_shape = output_shape
+    else
+        allocate(res % output_shape(size(output_shape)))
+        res % output_shape = output_shape
+    end if
   end function reshape_layer_cons
 
+
   pure module subroutine backward(self, input, gradient)
     class(reshape_generalized_layer), intent(in out) :: self
     real, intent(in) :: input(:)
@@ -56,14 +64,22 @@ module subroutine init(self, input_shape)
 
     self % input_shape = input_shape
 
-    ! Allocate gradient buffer based on input size
-    allocate(self % gradient(product(input_shape)))
+    !! Handle scalar input (size 1) or non-scalar
+    if (size(input_shape) == 1) then
+        allocate(self % gradient(1))
+    else
+        allocate(self % gradient(product(input_shape)))
+    end if
     self % gradient = 0
 
-    ! Allocate output buffer based on output_shape
-    allocate(self % output(product(self % output_shape)))
+    !! Handle scalar output_shape (size 1) or non-scalar
+    if (size(self % output_shape) == 1) then
+        allocate(self % output(1))
+    else
+        allocate(self % output(product(self % output_shape)))
+    end if
     self % output = 0
-  
-  end subroutine init
+end subroutine init
+
 
 end submodule nf_reshape_layer_generalized_submodule