get_partial_gno_kernel_coords_val Subroutine

  pure subroutine get_partial_gno_kernel_coords_val( &
       this, upstream_grad, output)
    !! In-place gradient w.r.t. edge features
    !!
    !! Chain rule through kernel:
    !!   kappa_e = V @ relu(U @ dx_e + b_u) + b_v
    !!   d(kappa_e)/d(dx_e) = V @ diag(relu'(U dx_e + b_u)) @ U
    !! Since the left operand already stores edge features directly,
    !! gradients accumulate independently for each edge column.
    implicit none

    ! Arguments
    class(array_type), intent(in) :: this
    !! Forward result node containing saved operands
    real(real32), dimension(:,:), intent(in)  :: upstream_grad
    !! Upstream gradient values
    real(real32), dimension(:,:), intent(out) :: output
    !! Output gradient values for coordinates

    ! Local variables
    integer :: d, H, F, num_e, e, k
    !! Unpacked dimensions and loop indices
    integer :: off_U, off_bu, off_V
    !! Flat offsets for packed kernel parameter blocks
    real(real32), allocatable :: U(:,:), b_u(:), V(:,:)
    !! Unpacked kernel parameter tensors
    real(real32), allocatable :: dx(:), pre_act(:), relu_mask(:)
    !! Per-edge buffers for input, pre-activation and ReLU mask
    real(real32), allocatable :: dkappa_ddx(:,:)   ! [F, d]
    !! Jacobian of edge kernel values with respect to coordinates
    real(real32), allocatable :: grad_dx(:)         ! [d]
    !! Coordinate gradient for one edge

    d = this%indices(1)
    H = this%indices(2)
    F = this%indices(3) * this%indices(4)
    num_e = size(this%left_operand%val, 2)

    off_U  = 0
    off_bu = H * d
    off_V  = off_bu + H

    allocate(U(H, d))
    U = reshape(this%right_operand%val(off_U+1:off_bu, 1), [H, d])
    allocate(b_u(H))
    b_u = this%right_operand%val(off_bu+1:off_V, 1)
    allocate(V(F, H))
    V = reshape(this%right_operand%val(off_V+1:off_V+F*H, 1), [F, H])

    allocate(dx(d), pre_act(H), relu_mask(H))
    allocate(dkappa_ddx(F, d), grad_dx(d))

    output = 0.0_real32

    do e = 1, num_e
       dx = this%left_operand%val(:, e)
       pre_act = matmul(U, dx) + b_u
       do k = 1, H
          if(pre_act(k) .gt. 0.0_real32)then
             relu_mask(k) = 1.0_real32
          else
             relu_mask(k) = 0.0_real32
          end if
       end do

       dkappa_ddx = 0.0_real32
       do k = 1, H
          if(relu_mask(k) .gt. 0.0_real32)then
             dkappa_ddx = dkappa_ddx + &
                  spread(V(:, k), 2, d) * spread(U(k, :), 1, F)
          end if
       end do

       grad_dx = matmul(upstream_grad(:, e), dkappa_ddx)
       output(:, e) = output(:, e) + grad_dx
    end do

    deallocate(U, b_u, V, dx, pre_act, relu_mask, dkappa_ddx, grad_dx)

  end subroutine get_partial_gno_kernel_coords_val