spread_array Module Function

  module function spread_array(source, dim, index, ncopies) result(c)
    !! Spread an autodiff array along a dimension
    implicit none
    class(array_type), intent(in), target :: source
    integer, intent(in) :: dim
    integer, intent(in) :: index
    integer, intent(in) :: ncopies
    type(array_type), pointer :: c

    integer :: s

    if(size(source%shape) .ne. 1)then
       call stop_program("spread: only 1D arrays can be used")
    end if

    if(dim.eq.1)then
       c => source%create_result(array_shape=[ncopies, size(source%val,2)])
       do concurrent(s=1:ncopies)
          c%val(s, :) = source%val(index, :)
       end do
    else if(dim.eq.2)then
       c => source%create_result(array_shape=[size(source%val,1), ncopies])
       do concurrent(s=1:ncopies)
          c%val(:, s) = source%val(:, index)
       end do
    else
       call stop_program("spread: only 1 or 2 dimensions are supported")
    end if
    c%indices = [index]
    allocate(c%adj_ja(2,1))
    c%adj_ja(:,1) = [ dim, size(source%val,dim) ]

    c%get_partial_left => get_partial_spread
    if(source%requires_grad)then
       c%requires_grad = .true.
       c%is_forward = source%is_forward
       c%operation = 'spread'
       c%left_operand => source
       c%owns_left_operand = source%is_temporary
    end if
  end function spread_array