get_partial_conv2d_kernel_val Subroutine

  pure subroutine get_partial_conv2d_kernel_val(this, upstream_grad, output)
    !! Get partial derivative wrt kernel for 2D convolution (subroutine version)
    implicit none

    class(array_type), intent(in) :: this
    real(real32), dimension(:,:), intent(in) :: upstream_grad
    real(real32), dimension(:,:), intent(out) :: output

    ! Local variables
    integer :: i, j, ki, kj, c_in, c_out, s
    integer :: i_in, j_in, k_idx, out_idx, in_idx
    integer :: in_base_idx, out_base_idx, k_base_idx, kernel_channel_size
    integer :: input_h, input_w, kernel_h, kernel_w
    integer :: output_h, output_w, num_channels, num_filters
    integer, dimension(2) :: stride, dilation
    integer :: channel_size_in, channel_size_out
    real(real32) :: grad_sum


    ! Unpack parameters
    num_channels = this%indices(1)
    num_filters = this%indices(2)
    stride = this%adj_ja(1:2,1)
    dilation = this%adj_ja(1:2,2)
    kernel_h = this%adj_ja(1,3)
    kernel_w = this%adj_ja(2,3)

    input_h = this%left_operand%shape(1)
    input_w = this%left_operand%shape(2)
    output_h = this%shape(1)
    output_w = this%shape(2)
    channel_size_in  = input_h * input_w
    channel_size_out = output_h * output_w
    kernel_channel_size = kernel_h * kernel_w

    output = 0._real32

    ! Parallelised over filters, channels, and kernel dimensions
    do concurrent(c_out = 1:num_filters, c_in = 1:num_channels, &
         kj = 1:kernel_w, ki = 1:kernel_h)
       out_base_idx = (c_out - 1) * channel_size_out
       in_base_idx = (c_in - 1) * channel_size_in
       k_base_idx = (c_in - 1) * kernel_channel_size + &
            (c_out - 1) * kernel_channel_size * num_channels
       k_idx = ki + (kj - 1) * kernel_h + k_base_idx

       grad_sum = 0._real32
       do s = 1, size(upstream_grad, dim=2)
          do j = 1, output_w
             j_in = (j - 1) * stride(2) + (kj - 1) * dilation(2) + 1
             if(j_in .ge. 1 .and. j_in .le. input_w)then
                do i = 1, output_h
                   i_in = (i - 1) * stride(1) + (ki - 1) * dilation(1) + 1
                   if(i_in .ge. 1 .and. i_in .le. input_h)then
                      in_idx  = i_in + (j_in - 1) * input_h + in_base_idx
                      out_idx = i + (j - 1) * output_h + out_base_idx
                      grad_sum = grad_sum + &
                           upstream_grad(out_idx, s) * this%left_operand%val(in_idx, s)
                   end if
                end do
             end if
          end do
       end do
       output(k_idx, 1) = grad_sum
    end do

  end subroutine get_partial_conv2d_kernel_val