update_message_gno Subroutine

  subroutine update_message_gno(this, input)
    !! Update message for the Graph Neural Operator layer
    !!
    !! Builds a differentiable computation graph through the diffstruc
    !! autodiff framework:
    !!
    !! input(1,s) : node features                    [F_in x num_vertices]
    !! input(2,s) : edge geometry / relative coords [coord_dim x num_edges]
    !!
    !! Pointer chain per sample s:
    !!   1. edge_kernels = gno_kernel_eval(coords, kernel_params, adj)
    !!   2. agg          = gno_aggregate(features, edge_kernels, adj)
    !!   3. bypass       = matmul(W, features)
    !!   4. z            = agg + bypass
    !!   5. z            = add_bias(z, b)          (if use_bias)
    !!   6. output       = activation(z)
    implicit none

    ! Arguments
    class(graph_nop_layer_type), intent(inout), target :: this
    !! Layer instance to execute
    class(array_type), dimension(:,:), intent(in), target :: input
    !! Input node-feature and edge-feature tensors

    ! Local variables
    integer :: s, F_in, F_out
    !! Sample index and input/output feature counts
    type(array_type), pointer :: ptr1, ptr2, ptr3, ptr4
    !! Intermediate kernel, aggregate, bypass and combined tensors

    F_in  = this%num_vertex_features(0)
    F_out = this%num_vertex_features(1)

    ! Allocate output array
    if(size(input, 1) .lt. 2)then
       call stop_program( &
            'graph_nop layer expects vertex and edge feature inputs' &
       )
       return
    end if

    if(allocated(this%output))then
       if(any(shape(this%output).ne.[2, size(input,2)]))then
          deallocate(this%output)
          allocate(this%output(2, size(input,2)))
       end if
    else
       allocate(this%output(2, size(input,2)))
    end if

    do s = 1, size(input, 2)

       ! Step 1: Evaluate kernel MLP on every edge
       ptr1 => gno_kernel_eval( &
            input(2,s), &               ! edge features [d, num_edges]
            this%params(1), &            ! packed kernel params
            this%graph(s)%adj_ia, &
            this%graph(s)%adj_ja, &
            this%coord_dim, this%kernel_hidden, F_in, F_out &
       )

       ! Step 2: Aggregate neighbour messages using per-edge kernels
       ptr2 => gno_aggregate( &
            input(1,s), &               ! features [F_in, num_v]
            ptr1, &                      ! edge kernels [F*F_in, num_edges]
            this%graph(s)%adj_ia, &
            this%graph(s)%adj_ja, &
            F_in, F_out &
       )

       ! Step 3: Linear bypass — W @ features
       ptr3 => matmul(this%params(2), input(1,s))

       ! Step 4: Combine aggregation with bypass
       ptr4 => ptr2 + ptr3

       ! Step 5: Add bias (if used)
       if(this%use_bias)then
          ptr4 => add_bias( &
               ptr4, this%params(3), dim=1, dim_act_on_shape=.true. &
          )
       end if

       ! Step 6: Apply activation
       ptr4 => this%activation%apply(ptr4)

       ! Store output
       call this%output(1,s)%zero_grad()
       call this%output(1,s)%assign_and_deallocate_source(ptr4)
       this%output(1,s)%is_temporary = .false.

       if(this%output(2,s)%allocated) call this%output(2,s)%deallocate()
       call this%output(2,s)%allocate(source=input(2,s)%val)
       call this%output(2,s)%zero_grad()
       call this%output(2,s)%set_requires_grad(.false.)
       this%output(2,s)%is_temporary = .false.
    end do

  end subroutine update_message_gno