concat_arrays Module Function

  module function concat_arrays(a, b, dim) result(c)
    !! Concatenate two autodiff arrays along the first dimension
    implicit none
    class(array_type), intent(in), target :: a, b
    integer, intent(in), optional :: dim
    type(array_type), pointer :: c

    integer :: i, s
    integer :: dim_
    integer, dimension(:), allocatable :: input_shape

    if(present(dim)) then
       dim_ = dim
    else
       dim_ = 1
    end if

    ! concatenate 1D array by using shape to swap dimensions
    if(dim_.eq.1)then
       if(size(a%val,2).ne.size(b%val,2)) then
          call stop_program( &
               "concat: arrays must have the same size in non-concat dimension" &
          )
          return
       elseif(size(a%shape).ne.size(b%shape)) then
          call stop_program( &
               "concat: arrays must have the same number of dimensions" &
          )
          return
       end if
       input_shape = a%shape
       input_shape(size(input_shape)) = &
            a%shape(size(input_shape)) + b%shape(size(input_shape))
       c => a%create_result(array_shape = &
            [input_shape, size(a%val,2)])
       c%val = 0._real32
       do concurrent(s=1:size(a%val,2))
          do concurrent(i=1:size(a%val,1))
             c%val(i, s) = a%val(i, s)
          end do
          do concurrent(i=1:size(b%val,1))
             c%val( size(a%val,1) + i, s) = b%val( i, s)
          end do
       end do
    elseif(dim_.eq.2)then
       if(size(a%val,1).ne.size(b%val,1)) then
          call stop_program( &
               "concat: arrays must have the same size in non-concat dimension" &
          )
          return
       end if
       c => a%create_result(array_shape = &
            [a%shape, size(a%val,2) + size(b%val,2)])
       c%val = 0._real32
       do concurrent(s=1:size(a%val,1))
          do concurrent(i=1:size(a%val,2))
             c%val(s, i) = a%val(s, i)
          end do
          do concurrent(i=1:size(b%val,2))
             c%val( s, size(a%val,2) + i) = b%val( s, i)
          end do
       end do
    else
       call stop_program("concat: only 1 or 2 dimensions are supported")
       return
    end if
    c%indices = [ dim_ ]

    c%get_partial_left => get_partial_concat_left
    c%get_partial_right => get_partial_concat_right
    c%get_partial_left_val => get_partial_concat_left_val
    c%get_partial_right_val => get_partial_concat_right_val
    if(a%requires_grad .or. b%requires_grad) then
       c%requires_grad = .true.
       c%is_forward = a%is_forward .or. b%is_forward
       c%operation = 'concat'
       c%left_operand => a
       c%right_operand => b
       c%owns_left_operand = a%is_temporary
       c%owns_right_operand = b%is_temporary
    end if
  end function concat_arrays