set_gradients Module Subroutine

  module subroutine set_gradients(this, gradients)
    !! Set the gradients of the layer
    !!
    !! This function sets the gradients of the layer from a single array.
    !! This has been modified from the neural-fortran library
    implicit none

    ! Arguments
    class(learnable_layer_type), intent(inout) :: this
    !! Instance of the layer
    real(real32), dimension(..), intent(in) :: gradients
    !! Gradients of the layer

    ! Local variables
    integer :: i, start_idx, end_idx
    !! Loop indices

    start_idx = 0
    end_idx = 0
    select rank(gradients)
    rank(0)
       do i = 1, size(this%params)
          if(.not.associated(this%params(i)%grad))then
             this%params(i)%grad => this%params(i)%create_result()
          end if
          this%params(i)%grad%val(:,1) = gradients
       end do
    rank(1)
       do i = 1, size(this%params)
          if(.not.associated(this%params(i)%grad))then
             this%params(i)%grad => this%params(i)%create_result()
          end if
          start_idx = end_idx + 1
          end_idx = start_idx + size(this%params(i)%val,1) - 1
          this%params(i)%grad%val(:,1) = gradients(start_idx:end_idx)
       end do
    end select

  end subroutine set_gradients