get_output_shape Module Function

  module function get_output_shape(this) result(output_shape)
    !! Get the output of the network
    implicit none

    ! Arguments
    class(network_type), intent(in) :: this
    !! Instance of network
    integer, dimension(2) :: output_shape
    !! Output shape

    ! Local variables
    integer :: i, layer_idx
    !! Loop indices


    ! array data: [ layer idx, empty ]
    ! graph data: [ vertex/edge idx, sample idx]

    if(this%use_graph_output)then
       output_shape = [2, this%batch_size]
       do i = 1, size(this%leaf_vertices,1), 1
          layer_idx = this%auto_graph%vertex(this%leaf_vertices(i))%id
          if(size(this%model(layer_idx)%layer%output,2).ne.this%batch_size)then
             call stop_program( &
                  "Inconsistent batch size in output layers" &
             )
             return
          end if
          output_shape(1) = output_shape(1) + &
               size( this%model(layer_idx)%layer%output, 1 )
       end do
    else
       output_shape = [0, 1]
       do i = 1, size(this%leaf_vertices,1)
          layer_idx = this%auto_graph%vertex(this%leaf_vertices(i))%id
          if(size(this%model(layer_idx)%layer%output,2).ne.1)then
             call stop_program( &
                  "Inconsistent size of dimension 2 in output layers" &
             )
             return
          end if
          output_shape(1) = &
               output_shape(1) + size( this%model(layer_idx)%layer%output, 1 )
       end do
    end if

  end function get_output_shape