duvenaud_propagate Module Function

  module function duvenaud_propagate( &
       vertex_features, edge_features, adj_ia, adj_ja &
  ) result(c)
    !! Propagate values from one autodiff array to another
    implicit none

    ! Arguments
    class(array_type), intent(in), target :: vertex_features, edge_features
    !! Vertex and edge feature tensors
    integer, dimension(:), intent(in) :: adj_ia
    !! CSR row pointers
    integer, dimension(:,:), intent(in) :: adj_ja
    !! CSR neighbour and edge lookup indices
    type(array_type), pointer :: c
    !! Propagated concatenated feature tensor

    ! Local variables
    integer :: v, w
    !! Vertex and adjacency traversal indices

    c => vertex_features%create_result( &
         array_shape = [ &
              size(vertex_features%val,1) + size(edge_features%val,1), &
              size(vertex_features%val,2) &
         ] &
    )
    ! propagate 1D array by using shape to swap dimensions
    do concurrent(v=1:size(vertex_features%val,2))
       c%val(:,v) = 0.0_real32
       do w = adj_ia(v), adj_ia(v+1)-1
          c%val(:,v) = c%val(:,v) + [ &
               vertex_features%val(:, adj_ja(1, w)), &
               edge_features%val(:, adj_ja(2, w)) &
          ]
       end do
    end do

    c%indices = adj_ia
    c%adj_ja = adj_ja
    c%get_partial_left => get_partial_duvenaud_propagate_left
    c%get_partial_right => get_partial_duvenaud_propagate_right
    c%get_partial_left_val => get_partial_duvenaud_propagate_left_val
    c%get_partial_right_val => get_partial_duvenaud_propagate_right_val
    if(vertex_features%requires_grad .or. edge_features%requires_grad)then
       c%requires_grad = .true.
       c%is_forward = vertex_features%is_forward .or. edge_features%is_forward
       c%operation = 'duvenaud_propagate'
       c%left_operand => vertex_features
       c%right_operand => edge_features
       c%owns_left_operand = vertex_features%is_temporary
       c%owns_right_operand = edge_features%is_temporary
    end if
  end function duvenaud_propagate