real2d_matmul Module Function

  module function real2d_matmul(a, b) result(c)
    !! Matrix multiplication of a real 2D matrix and an autodiff array
    !! Computes C = a * b where a is a real matrix and b is autodiff
    implicit none
    real(real32), dimension(:,:), intent(in) :: a
    class(array_type), intent(in), target :: b
    type(array_type), pointer :: c
    type(array_type), pointer :: a_array

    integer :: s, i, m, k, num_samples

    m = size(a, 1)
    k = size(a, 2)
    num_samples = size(b%val, 2)

    c => b%create_result(array_shape = [m, num_samples])
#ifdef USE_BLAS
    ! C(m, S) = a(m, k) * b%val(k, S)
    ! sgemm: standard C = alpha * A * B, no transposes
    ! m = m, n = S, k = k
    ! lda = m, ldb = k, ldc = m
    call sgemm('N', 'N', m, num_samples, k, &
         1.0_real32, a, m, b%val, k, 0.0_real32, c%val, m)
#else
    c%val = matmul(a, b%val)
#endif

    c%is_sample_dependent = b%is_sample_dependent
    c%get_partial_left => get_partial_matmul_left
    c%get_partial_right => get_partial_matmul_right
    c%get_partial_left_val => get_partial_matmul_left_val
    c%get_partial_right_val => get_partial_matmul_right_val
    c%get_partial_left_val_sum => get_partial_matmul_left_val_sum
    c%get_partial_right_val_sum => get_partial_matmul_right_val_sum
    if(b%requires_grad) then
       c%requires_grad = .true.
       c%is_forward = b%is_forward
       c%operation = 'matmul_scalar'
       c%right_operand => b
       c%owns_right_operand = b%is_temporary
    end if
    allocate(a_array)
    a_array%is_sample_dependent = .false.
    a_array%shape = shape(a)
    a_array%requires_grad = .false.
    call a_array%allocate(array_shape=[size(a,1), size(a,2), 1])
    a_array%val(:,1) = reshape(a, [size(a,1)*size(a,2)])
    c%left_operand => a_array
    c%owns_left_operand = .true.
  end function real2d_matmul