mean_array Module Function

  module function mean_array(a, dim) result(c)
    !! Compute mean values along a dimension - optimized version
    implicit none
    class(array_type), intent(in), target :: a
    integer, intent(in) :: dim
    type(array_type), pointer :: c

    integer :: s, i, n_rows, n_cols
    real(real32) :: rtmp1, inv_count

    ! if(size(a%shape) .ne. 1)then
    !    call stop_program("mean_array: only 1D arrays can be used")
    ! end if

    ! Cache dimensions to avoid repeated size() calls
    n_rows = size(a%val, 1)
    n_cols = size(a%val, 2)

    if(dim.eq.1)then
       c => a%create_result(array_shape = [1, n_cols])
       rtmp1 = real(n_rows, real32)
       inv_count = 1.0_real32 / rtmp1

       c%val(1,:) = sum(a%val, dim=1) * inv_count

    else if(dim.eq.2)then
       c => a%create_result(array_shape = [a%shape, 1])
       rtmp1 = real(n_cols, real32)
       inv_count = 1.0_real32 / rtmp1

       c%val(:,1) = sum(a%val, dim=2) * inv_count

       c%is_sample_dependent = .false.
    else
       call stop_program("mean_array: only 1 or 2 dimensions are supported")
    end if
    c%indices = [dim, 1]

    c%get_partial_left => get_partial_mean
    c%get_partial_left_val => get_partial_mean_val
    if(a%requires_grad)then
       c%requires_grad = .true.
       c%is_forward = a%is_forward
       c%operation = 'mean_array'
       c%left_operand => a
       c%owns_left_operand = a%is_temporary
    end if

  end function mean_array