accumulate_edge_gradients_2d_val Subroutine

  pure subroutine accumulate_edge_gradients_2d_val(upstream_grad, output, &
       input_shape, indices, adj_ja)
    !! Accumulate edge gradients for 2D padding - raw array version
    implicit none

    ! Arguments
    real(real32), dimension(:,:), intent(in) :: upstream_grad
    real(real32), dimension(:,:), intent(inout) :: output
    integer, dimension(4), intent(in) :: input_shape
    integer, dimension(:), intent(in) :: indices
    integer, dimension(:,:), intent(in) :: adj_ja

    ! Local variables
    integer :: i, j, m, s, f, idx
    integer :: idx_in, idx_out
    integer :: input_size_h, input_size_w
    integer :: output_size_h, output_size_w
    integer :: num_edge_facets
    integer :: facet_dim
    real(real32) :: grad_sum

    input_size_h = input_shape(1)
    input_size_w = input_shape(2)
    output_size_h = input_size_h + 2 * indices(2)
    output_size_w = input_size_w + 2 * indices(3)
    num_edge_facets = indices(4)

    if(num_edge_facets .eq. 0) return

    select case(indices(1))
    case(3) ! circular
       do f = 1, num_edge_facets
          facet_dim = indices(5 + f)
          if(facet_dim .eq. 1)then
             do concurrent( &
                  s = 1:input_shape(4), &
                  m = 1:input_shape(3), &
                  j = adj_ja(1,(f-1)*4 + 3):adj_ja(1,(f-1)*4 + 4), &
                  i = adj_ja(1,(f-1)*4 + 1):adj_ja(1,(f-1)*4 + 2))
                idx_in = i + (j-1) * input_size_h + &
                     (m-1) * input_size_h * input_size_w
                idx_out = &
                     ( &
                          adj_ja(2,(f-1)*4 + 1) + &
                          (i - adj_ja(1,(f-1)*4 + 1)) &
                     ) + &
                     (j + adj_ja(2,(f-1)*4 + 3) - adj_ja(1,(f-1)*4 + 3) - 1) * &
                     output_size_h + (m-1) * output_size_h * output_size_w
                output(idx_in, s) = output(idx_in, s) + &
                     upstream_grad(idx_out, s)
             end do
          else
             do concurrent( &
                  s = 1:input_shape(4), &
                  m = 1:input_shape(3), &
                  j = adj_ja(1,(f-1)*4 + 3):adj_ja(1,(f-1)*4 + 4), &
                  i = adj_ja(1,(f-1)*4 + 1):adj_ja(1,(f-1)*4 + 2))
                idx_in = i + (j-1) * input_size_h + &
                     (m-1) * input_size_h * input_size_w
                idx_out = &
                     ( &
                          i + adj_ja(2,(f-1)*4 + 1) - &
                          adj_ja(1,(f-1)*4 + 1) &
                     ) + ( &
                          adj_ja(2,(f-1)*4 + 3) + &
                          (j - adj_ja(1,(f-1)*4 + 3)) - 1 &
                     ) * output_size_h + &
                     (m-1) * output_size_h * output_size_w
                output(idx_in, s) = output(idx_in, s) + &
                     upstream_grad(idx_out, s)
             end do
          end if
       end do
    case(4) ! reflection
       do f = 1, num_edge_facets
          facet_dim = indices(5 + f)
          if(facet_dim .eq. 1)then
             do concurrent( &
                  s = 1:input_shape(4), &
                  m = 1:input_shape(3), &
                  j = adj_ja(1,(f-1)*4 + 3):adj_ja(1,(f-1)*4 + 4), &
                  i = adj_ja(1,(f-1)*4 + 2):adj_ja(1,(f-1)*4 + 1))
                idx_in = i + (j-1) * input_size_h + &
                     (m-1) * input_size_h * input_size_w
                idx_out = &
                     ( &
                          adj_ja(2,(f-1)*4 + 1) - &
                          (i - adj_ja(1,(f-1)*4 + 1)) &
                     ) + &
                     (j + adj_ja(2,(f-1)*4 + 3) - adj_ja(1,(f-1)*4 + 3) - 1) * &
                     output_size_h + (m-1) * output_size_h * output_size_w
                output(idx_in, s) = output(idx_in, s) + &
                     upstream_grad(idx_out, s)
             end do
          else
             do concurrent( &
                  s = 1:input_shape(4), &
                  m = 1:input_shape(3), &
                  j = adj_ja(1,(f-1)*4 + 4):adj_ja(1,(f-1)*4 + 3), &
                  i = adj_ja(1,(f-1)*4 + 1):adj_ja(1,(f-1)*4 + 2))
                idx_in = i + (j-1) * input_size_h + &
                     (m-1) * input_size_h * input_size_w
                idx_out = &
                     ( &
                          i + adj_ja(2,(f-1)*4 + 1) - &
                          adj_ja(1,(f-1)*4 + 1) &
                     ) + ( &
                          adj_ja(2,(f-1)*4 + 4) - &
                          (j - adj_ja(1,(f-1)*4 + 4)) - 1 &
                     ) * output_size_h + &
                     (m-1) * output_size_h * output_size_w
                output(idx_in, s) = output(idx_in, s) + &
                     upstream_grad(idx_out, s)
             end do
          end if
       end do
    case(5) ! replication
       do f = 1, num_edge_facets
          facet_dim = indices(5 + f)
          if(facet_dim .eq. 1)then
             do s = 1, input_shape(4)
                do m = 1, input_shape(3)
                   do j = adj_ja(1,(f-1)*4 + 3), adj_ja(1,(f-1)*4 + 4)
                      grad_sum = 0._real32
                      do i = adj_ja(2,(f-1)*4 + 1), adj_ja(2,(f-1)*4 + 2)
                         idx_out = i + &
                              ( &
                                   j + adj_ja(2,(f-1)*4 + 3) - &
                                   adj_ja(1,(f-1)*4 + 3) - 1 &
                              ) * output_size_h + &
                              (m-1) * output_size_h * output_size_w
                         grad_sum = grad_sum + upstream_grad(idx_out, s)
                      end do
                      idx_in = adj_ja(1,(f-1)*4 + 1) + (j-1) * input_size_h + &
                           (m-1) * input_size_h * input_size_w
                      output(idx_in, s) = output(idx_in, s) + grad_sum
                   end do
                end do
             end do
          else
             do s = 1, input_shape(4)
                do m = 1, input_shape(3)
                   do i = adj_ja(1,(f-1)*4 + 1), adj_ja(1,(f-1)*4 + 2)
                      grad_sum = 0._real32
                      do j = adj_ja(2,(f-1)*4 + 3), adj_ja(2,(f-1)*4 + 4)
                         idx_out = &
                              ( i + adj_ja(2,(f-1)*4 + 1) - adj_ja(1,(f-1)*4 + 1) ) + &
                              (j-1) * output_size_h + &
                              (m-1) * output_size_h * output_size_w
                         grad_sum = grad_sum + upstream_grad(idx_out, s)
                      end do
                      idx_in = i + (adj_ja(1,(f-1)*4 + 3) - 1) * &
                           input_size_h + (m-1) * input_size_h * input_size_w
                      output(idx_in, s) = output(idx_in, s) + grad_sum
                   end do
                end do
             end do
          end if
       end do
    end select

  end subroutine accumulate_edge_gradients_2d_val