unpack_mask_array Module Function

  module function unpack_mask_array(a, array_shape, dim, indices) result(c)
    !! Unpack an autodiff array
    implicit none
    class(array_type), intent(in), target :: a
    integer, dimension(:), intent(in) :: array_shape
    integer, intent(in), optional :: dim
    integer, dimension(:), intent(in), optional :: indices
    type(array_type), pointer :: c

    integer :: i, s, dim_, num_samples, num_elements

    if(present(dim)) then
       dim_ = dim
    else
       dim_ = 1
    end if
    num_samples = array_shape(size(array_shape))
    num_elements = product(array_shape(1:size(array_shape)-1))

    c => a%create_result(array_shape = array_shape)
    c%val = 0.0_real32
    if(dim_.eq.1)then
       if(present(indices))then
          do concurrent(i=1:size(indices,1), s=1:num_samples)
             c%val(indices(i),s) = a%val(i,s)
          end do
       else
          do concurrent(i=1:num_elements, s=1:num_samples)
             c%val(i,s) = a%val(i,s)
          end do
       end if
    elseif(dim_.eq.2)then
       if(present(indices))then
          do concurrent(i=1:num_elements, s=1:size(indices,1))
             c%val(i,indices(s)) = a%val(i,s)
          end do
       else
          do concurrent(i=1:num_elements, s=1:num_samples)
             c%val(i,s) = a%val(i,s)
          end do
       end if
    else
       call stop_program("unpack_mask: only 1 or 2 dimensions are supported")
    end if

    if(present(indices))then
       c%indices = indices
    end if
    allocate(c%adj_ja(1,1))
    c%adj_ja(1,1) = dim_

    c%get_partial_left => get_partial_unpack_mask
    c%get_partial_right => get_partial_pack_mask
    if(a%requires_grad) then
       c%requires_grad = .true.
       c%is_forward = a%is_forward
       c%operation = 'unpack_mask'
       c%left_operand => a
       c%owns_left_operand = a%is_temporary
    end if
  end function unpack_mask_array