conv1d Module Function

  module function conv1d(input, kernel, stride, dilation) result(output)
    !! 1D convolution operation
    implicit none

    ! Arguments
    type(array_type), intent(in), target :: input
    type(array_type), intent(in), target :: kernel
    integer, intent(in) :: stride
    integer, intent(in) :: dilation
    type(array_type), pointer :: output

    ! Local variables
    integer :: i, k, c_in, c_out, s
    integer :: i_in, k_idx
    integer :: input_h, kernel_h, output_h, num_channels, num_filters
    real(real32) :: conv_sum
    integer, dimension(3) :: output_shape

    ! Extract dimensions
    ! input: [H_in, C_in, B]
    ! kernel: [K, C_in, C_out]
    input_h = input%shape(1)
    num_channels = input%shape(2)
    kernel_h = kernel%shape(1)
    num_filters = kernel%shape(3)

    ! Calculate output dimensions
    output_h = (input_h - dilation*(kernel_h - 1) - 1) / &
         stride + 1
    output_shape = [output_h, num_filters, size(input%val, dim=2)]

    output => input%create_result(array_shape = output_shape)
    output%val = 0._real32

    ! Perform convolution
    do concurrent(s = 1:output_shape(3), c_out = 1:num_filters, &
         i = 1:output_h)
       conv_sum = 0._real32
       do c_in = 1, num_channels
          do k = 1, kernel_h
             i_in = ( i - 1 ) * stride + ( k - 1 ) * dilation + 1
             if(i_in .ge. 1 .and. i_in .le. input_h)then
                k_idx = k + ( c_in - 1 ) * kernel_h + &
                     ( c_out - 1 ) * kernel_h * num_channels
                conv_sum = conv_sum + &
                     input%val(i_in + ( c_in - 1 ) * input_h, s) * &
                     kernel%val(k_idx, 1)
             end if
          end do
       end do
       output%val(i + (c_out-1)*output_h, s) = conv_sum
    end do

    ! Store parameters for backward pass
    allocate(output%indices(2))
    output%indices(1) = num_channels
    output%indices(2) = num_filters
    allocate(output%adj_ja(1,3))
    output%adj_ja(1,1) = stride
    output%adj_ja(1,2) = dilation
    output%adj_ja(1,3) = kernel_h

    output%get_partial_left => get_partial_conv1d_input
    output%get_partial_right => get_partial_conv1d_kernel
    output%get_partial_left_val => get_partial_conv1d_input_val
    output%get_partial_right_val => get_partial_conv1d_kernel_val
    if(input%requires_grad .or. kernel%requires_grad)then
       output%requires_grad = .true.
       output%is_forward = input%is_forward
       output%operation = 'conv1d'
       output%left_operand => input
       output%right_operand => kernel
    end if

  end function conv1d