sum_and_pad_array Module Function

module function sum_and_pad_array(a, dim, new_dim_index, new_dim_size) result(c)

Sum values along a dimension and return an autodiff array

Arguments

Type IntentOptional Attributes Name
class(array_type), intent(in), target :: a
integer, intent(in) :: dim
integer, intent(in) :: new_dim_index
integer, intent(in) :: new_dim_size

Return Value type(array_type), pointer


Source Code

  module function sum_and_pad_array(a, dim, new_dim_index, new_dim_size) result(c)
    !! Sum values along a dimension and return an autodiff array
    implicit none
    class(array_type), intent(in), target :: a
    integer, intent(in) :: dim
    integer, intent(in) :: new_dim_index
    integer, intent(in) :: new_dim_size
    type(array_type), pointer :: c

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

    ! sum 1D array by using shape to swap dimensions
    if(dim.eq.1)then
       c => a%create_result(array_shape=[new_dim_size, size(a%val,2)])
       c%val = 0.0_real32
       c%val(new_dim_index,:) = sum(a%val(:,:), dim=1)
    else if(dim.eq.2)then
       c => a%create_result(array_shape=[size(a%val,1), new_dim_size])
       c%val = 0.0_real32
       c%val(:,new_dim_index) = sum(a%val(:,:), dim=2)
    end if
    c%indices = [dim, new_dim_index]

    c%get_partial_left => get_partial_sum_and_pad
    c%get_partial_left_val => get_partial_sum_and_pad_val
    c%is_sample_dependent = a%is_sample_dependent
    if(a%requires_grad)then
       c%requires_grad = .true.
       c%is_forward = a%is_forward
       c%operation = 'sum_and_pad_array'
       c%left_operand => a
       c%owns_left_operand = a%is_temporary
    end if
    c%indices = [dim, new_dim_index]
  end function sum_and_pad_array