get_partial_matmul_right_val_sum Subroutine

  subroutine get_partial_matmul_right_val_sum(this, upstream_grad, output)
#else
  pure subroutine get_partial_matmul_right_val_sum(this, upstream_grad, output)
#endif
    !! Sum-reduced gradient w.r.t. right operand for matmul.
    !! For the outer product case (rank-1 left operand), this computes:
    !!   output = sum_s(left(:,s) (x) upstream(:,s)) = left * upstream^T
    !! using a single SGEMM call.
    implicit none
    class(array_type), intent(in) :: this
    real(real32), dimension(:,:), intent(in) :: upstream_grad
    real(real32), dimension(:), intent(out) :: output

    integer :: i, j, s, m, n, num_elements, num_batches, num_upstream

    num_batches = size(upstream_grad, 2)

    if(size(this%left_operand%shape).eq.2)then
       ! 2D case: output = sum_s(W^T * upstream(:,s)) = W^T * sum(upstream, dim=2)
       m = this%left_operand%shape(1)
       n = this%left_operand%shape(2)
       if(this%left_operand%is_sample_dependent)then
          block
            real(real32), dimension(m, n) :: temp
            real(real32), dimension(n) :: col_result
            output = 0.0_real32
            do s = 1, num_batches
               temp = reshape(this%left_operand%val(:,s), [m, n])
               col_result = matmul(transpose(temp), upstream_grad(:,s))
               output = output + col_result
            end do
          end block
       else
#ifdef USE_BLAS
          block
            real(real32), pointer :: W(:,:)
            real(real32), dimension(size(upstream_grad,1)) :: upstream_sum
            W(1:m, 1:n) => this%left_operand%val(:,1)
            upstream_sum = sum(upstream_grad, dim=2)
            ! W^T * upstream_sum: sgemv with 'T'
            call sgemv('T', m, n, 1.0_real32, W, m, upstream_sum, 1, &
                 0.0_real32, output, 1)
          end block
#else
          block
            real(real32), dimension(n, m) :: temp_t
            real(real32), dimension(size(upstream_grad,1)) :: upstream_sum
            temp_t = transpose(reshape(this%left_operand%val(:,1), [m, n]))
            upstream_sum = sum(upstream_grad, dim=2)
            output = matmul(temp_t, upstream_sum)
          end block
#endif
       end if
    else
       ! Outer product case: sum_s(left(i,s) * upstream(j,s))
       ! = matmul(left, upstream^T) stored column-major
       num_elements = size(this%left_operand%val, 1)
       num_upstream = size(upstream_grad, 1)
       if(this%left_operand%is_sample_dependent)then
#ifdef USE_BLAS
          call sgemm('N', 'T', num_elements, num_upstream, num_batches, &
               1.0_real32, this%left_operand%val, num_elements, &
               upstream_grad, num_upstream, &
               0.0_real32, output, num_elements)
#else
          output = 0.0_real32
          do s = 1, num_batches
             do j = 1, num_upstream
                do i = 1, num_elements
                   output((j-1)*num_elements + i) = &
                        output((j-1)*num_elements + i) + &
                        upstream_grad(j,s) * this%left_operand%val(i,s)
                end do
             end do
          end do
#endif
       else
          block
            real(real32), dimension(num_upstream) :: upstream_sum
            upstream_sum = sum(upstream_grad, dim=2)
            do j = 1, num_upstream
               output((j-1)*num_elements+1:j*num_elements) = &
                    this%left_operand%val(:, 1) * upstream_sum(j)
            end do
          end block
       end if
    end if

  end subroutine get_partial_matmul_right_val_sum