get_partial_mean_val Subroutine

  pure subroutine get_partial_mean_val(this, upstream_grad, output)
    implicit none
    class(array_type), intent(in) :: this
    real(real32), dimension(:,:), intent(in) :: upstream_grad
    real(real32), dimension(:,:), intent(out) :: output

    integer :: i, s, dim, n_rows, n_cols
    real(real32) :: inv_count

    ! Cache values to avoid repeated accesses
    dim = this%indices(1)
    n_rows = size(output, 1)
    n_cols = size(output, 2)
    inv_count = 1.0_real32 / real(size(this%left_operand%val, dim), real32)

    if(dim.eq.1)then
       do concurrent(s = 1:n_cols, i = 1:n_rows)
          output(i,s) = upstream_grad(1,s) * inv_count
       end do
    else if(dim.eq.2)then
       do concurrent(s = 1:n_cols, i = 1:n_rows)
          output(i,s) = upstream_grad(i,1) * inv_count
       end do
    end if
  end subroutine get_partial_mean_val