update Module Subroutine

  module subroutine update(this)
    !! Update the network
    implicit none

    ! Arguments
    class(network_type), intent(inout) :: this
    !! Instance of network
    real(real32), dimension(this%num_params) :: params, gradients
    !! Parameters and gradients

    ! Local variables
    integer :: l, i, s, start_idx, end_idx, seg
    !! Loop index


    !---------------------------------------------------------------------------
    ! Increment optimiser iteration counter
    !---------------------------------------------------------------------------
    if(this%optimiser%lr_decay%iterate_per_epoch)then
       if(this%epoch.gt.this%optimiser%epoch)then
          this%optimiser%epoch = this%epoch
          this%optimiser%iter = this%optimiser%iter + 1
       end if
    else
       this%optimiser%iter = this%optimiser%iter + 1
    end if


    !---------------------------------------------------------------------------
    ! Get learnable parameters and gradients (using pre-computed layout)
    !---------------------------------------------------------------------------
    if(this%param_num_segments.gt.0 .and. &
         allocated(this%param_seg_layer))then
       do seg = 1, this%param_num_segments
          l = this%param_seg_layer(seg)
          i = this%param_seg_pidx(seg)
          start_idx = this%param_seg_start(seg)
          end_idx = this%param_seg_end(seg)
          select type(current => this%model(l)%layer)
          class is(learnable_layer_type)
             params(start_idx:end_idx) = current%params(i)%val(:,1)
             if(.not.associated(current%params(i)%grad))then
                call stop_program( &
                     "Gradient not allocated for parameters" &
                )
             end if
             s = size(current%params(i)%grad%val,2)
             if(s.eq.1)then
                gradients(start_idx:end_idx) = &
                     current%params(i)%grad%val(:,1)
             else
                gradients(start_idx:end_idx) = &
                     sum(current%params(i)%grad%val, dim=2) / &
                     real(s, real32)
             end if
          end select
       end do
    else
       start_idx = 0
       end_idx   = 0
       do l = 1, this%num_layers
          select type(current => this%model(l)%layer)
          class is(learnable_layer_type)
             do i = 1, size(current%params)
                start_idx = end_idx + 1
                end_idx = end_idx + size(current%params(i)%val, 1)
                params(start_idx:end_idx) = current%params(i)%val(:,1)
                if(.not.associated(current%params(i)%grad))then
                   call stop_program( &
                        "Gradient not allocated for parameters" &
                   )
                end if
                select case(size(current%params(i)%grad%val,2))
                case(1)
                   gradients(start_idx:end_idx) = &
                        current%params(i)%grad%val(:,1)
                case default
                   gradients(start_idx:end_idx) = [ &
                        sum(current%params(i)%grad%val, dim=2) / &
                        real( &
                             size(current%params(i)%grad%val, dim=2), &
                             real32) &
                   ]
                end select
             end do
          end select
       end do
    end if
    call this%optimiser%clip_dict%apply(size(gradients),gradients)

    !---------------------------------------------------------------------------
    ! Update layers of learnable layer types
    !---------------------------------------------------------------------------
    call this%optimiser%minimise(params, gradients)

    ! Set params back using pre-computed layout
    if(this%param_num_segments.gt.0 .and. &
         allocated(this%param_seg_layer))then
       do seg = 1, this%param_num_segments
          l = this%param_seg_layer(seg)
          i = this%param_seg_pidx(seg)
          start_idx = this%param_seg_start(seg)
          end_idx = this%param_seg_end(seg)
          select type(current => this%model(l)%layer)
          class is(learnable_layer_type)
             current%params(i)%val(:,1) = params(start_idx:end_idx)
          end select
       end do
    else
       call this%set_params(params)
    end if
    call this%reset_gradients()

  end subroutine update