get_partial_mean Function

function get_partial_mean(this, upstream_grad) result(output)

Arguments

Type IntentOptional Attributes Name
class(array_type), intent(inout) :: this
type(array_type), intent(in) :: upstream_grad

Return Value type(array_type)


Source Code

  function get_partial_mean(this, upstream_grad) result(output)
    implicit none
    class(array_type), intent(inout) :: this
    type(array_type), intent(in) :: upstream_grad
    type(array_type) :: output

    real(real32) :: rtmp1
    type(array_type), pointer :: ptr

    ! Calculate the number of elements that were averaged
    rtmp1 = real(size(this%left_operand%val, this%indices(1)), real32)

    if(this%is_forward)then
       ptr => sum( upstream_grad, dim = this%indices(1) ) / rtmp1
    else
       ptr => spread( &
            upstream_grad / rtmp1, &
            dim=this%indices(1), &
            index=this%indices(2), &
            ncopies= size(this%left_operand%val, this%indices(1)) &
       )
    end if
    call output%assign_and_deallocate_source(ptr)
  end function get_partial_mean