pack_mask_array Module Function

  module function pack_mask_array(a, dim, mask) result(c)
    !! Pack an autodiff array using a logical mask
    implicit none
    class(array_type), intent(in), target :: a
    integer, intent(in) :: dim
    logical, dimension(:), intent(in), optional :: mask
    type(array_type), pointer :: c

    integer :: i, j, s, itmp1

    if(present(mask))then
       itmp1 = count(mask)
       if(dim.eq.1)then
          c => a%create_result(array_shape=[itmp1, size(a%val,2)])
          allocate(c%indices(itmp1))
          i = 0
          do concurrent(s=1:size(a%val,2))
             do j = 1, size(a%val,1)
                if(mask(j)) then
                   i = i + 1
                   c%val(i, s) = a%val(j, s)
                end if
             end do
             i = 0
          end do
          c%indices = pack([(j, j=1,size(mask))], mask)
       elseif(dim.eq.2)then
          c => a%create_result(array_shape=[a%shape, itmp1])
          allocate(c%indices(itmp1))
          i = 0
          do concurrent(s=1:size(a%val,1))
             do j = 1, size(a%val,2)
                if(mask(j)) then
                   i = i + 1
                   c%val(s, i) = a%val(s, j)
                end if
             end do
             i = 0
          end do
          c%indices = pack([(j, j=1,size(mask))], mask)
       else
          call stop_program("pack_mask: only 1 or 2 dimensions are supported")
       end if
    else
       if(dim.eq.1)then
          c => a%create_result(array_shape=[size(a%val,1), size(a%val,2)])
          do concurrent(s=1:size(a%val,2), i=1:size(a%val,1))
             c%val(i, s) = a%val(i, s)
          end do
       elseif(dim.eq.2)then
          c => a%create_result()
          do concurrent(s=1:size(a%val,1), i=1:size(a%val,2))
             c%val(s, i) = a%val(s, i)
          end do
       else
          call stop_program("pack: only 1 or 2 dimensions are supported")
       end if
       !c%indices = [(j, j=1,size(a%val,dim))]
    end if
    allocate(c%adj_ja(size(a%shape)+2,1))
    c%adj_ja(1,1) = dim
    c%adj_ja(2:,1) = [ a%shape, size(a%val,2) ]

    c%get_partial_left => get_partial_pack_mask
    c%get_partial_left_val => get_partial_pack_mask_val
    if(a%requires_grad) then
       c%requires_grad = .true.
       c%is_forward = a%is_forward
       c%operation = 'pack_mask'
       c%left_operand => a
       c%owns_left_operand = a%is_temporary
    end if
  end function pack_mask_array