avgpool2d Module Function

  module function avgpool2d(input, pool_size, stride) result(output)
    !! 2D average pooling operation
    implicit none

    ! Arguments
    type(array_type), intent(in), target :: input
    integer, dimension(2), intent(in) :: pool_size
    integer, dimension(2), intent(in) :: stride
    type(array_type), pointer :: output

    ! Local variables
    integer :: i, j, m, s, i_step, j_step
    integer :: stride_idx, idx, multiplier
    integer :: channel_size_in, channel_size_out
    real(real32) :: pool_sum, pool_norm
    integer, dimension(4) :: output_shape

    output_shape = [ &
         (input%shape(1) - pool_size(1)) / stride(1) + 1, &
         (input%shape(2) - pool_size(2)) / stride(2) + 1, &
         input%shape(3), &
         size(input%val, dim=2)]
    output => input%create_result(array_shape = output_shape)
    pool_norm = 1.0_real32 / real(pool_size(1) * pool_size(2), real32)

    ! Pre-compute as integers
    channel_size_in = input%shape(1) * input%shape(2)
    channel_size_out = output_shape(1) * output_shape(2)

    do concurrent(&
         s = 1:output_shape(4), &
         m = 1:output_shape(3), &
         j = 1:output_shape(2), &
         i = 1:output_shape(1))

       ! Compute indices once
       stride_idx = (i-1)*stride(1) + &
            ((j-1)*stride(2)) * input%shape(1) + &
            (m-1) * channel_size_in
       idx = i + (j - 1) * output_shape(1) + (m - 1) * channel_size_out

       pool_sum = 0._real32
       do j_step = 0, pool_size(2)-1
          do i_step = 0, pool_size(1)-1
             pool_sum = pool_sum + &
                  input%val(stride_idx + i_step + j_step * input%shape(1) + 1, s)
          end do
       end do
       output%val(idx, s) = pool_sum * pool_norm
    end do
    allocate(output%adj_ja(2,2))
    output%adj_ja(:,1) = pool_size
    output%adj_ja(:,2) = stride

    output%get_partial_left => get_partial_avgpool2d
    output%get_partial_left_val => get_partial_avgpool2d_val
    if(input%requires_grad)then
       output%requires_grad = .true.
       output%is_forward = input%is_forward
       output%operation = 'avgpool'
       output%left_operand => input
    end if

  end function avgpool2d