unspread_array Module Function

  module function unspread_array(source, index, dim, new_size) result(c)
    !! Unpack an autodiff array
    implicit none
    class(array_type), intent(in), target :: source
    integer, intent(in) :: index
    integer, intent(in) :: new_size, dim
    type(array_type), pointer :: c

    integer :: i, s


    if(dim.eq.1)then
       c => source%create_result(array_shape = [ new_size, size(source%val,2) ])
       c%val = 0.0_real32
       do concurrent(i=1:size(source%val,1), s=1:size(source%val,2))
          c%val(index,s) = c%val(index,s) + source%val(i,s)
       end do
    elseif(dim.eq.2)then
       c => source%create_result( array_shape = [ size(source%val,1), new_size ] )
       c%val = 0.0_real32
       do concurrent(i=1:size(source%val,1), s=1:size(source%val,2))
          c%val(i,index) = c%val(i,index) + source%val(i,s)
       end do
    end if
    c%indices = [index]
    allocate(c%adj_ja(2,1))
    c%adj_ja(:,1) = [ dim, new_size ]

    c%get_partial_left => get_partial_unspread
    if(source%requires_grad)then
       c%requires_grad = .true.
       c%is_forward = source%is_forward
       c%operation = 'unspread'
       c%left_operand => source
       c%owns_left_operand = source%is_temporary
    end if
  end function unspread_array