fix: resolve compilation errors in weighted/generalized lstsq

Author: aamrindersingh

src/linalg/stdlib_linalg.fypp

@@ -703,8 +703,8 @@ module stdlib_linalg
     !!
     #:for rk,rt,ri in RC_KINDS_TYPES
     module function stdlib_linalg_${ri}$_weighted_lstsq(w,a,b,cond,overwrite_a,rank,err) result(x)
-        !> Weight vector (must be positive)
-        ${rt}$, intent(in) :: w(:)
+        !> Weight vector (must be positive, always real)
+        real(${rk}$), intent(in) :: w(:)
         !> Input matrix a[m,n]
         ${rt}$, intent(inout), target :: a(:,:)
         !> Right hand side vector b[m]

src/linalg/stdlib_linalg_least_squares.fypp

@@ -571,8 +571,8 @@ submodule (stdlib_linalg) stdlib_linalg_least_squares
     #:for rk,rt,ri in RC_KINDS_TYPES
     ! Weighted least-squares: minimize ||D(Ax - b)||^2 where D = diag(sqrt(w))
     module function stdlib_linalg_${ri}$_weighted_lstsq(w,a,b,cond,overwrite_a,rank,err) result(x)
-        !> Weight vector (must be positive)
-        ${rt}$, intent(in) :: w(:)
+        !> Weight vector (must be positive, always real)
+        real(${rk}$), intent(in) :: w(:)
         !> Input matrix a[m,n]
         ${rt}$, intent(inout), target :: a(:,:)
         !> Right hand side vector b[m]
@@ -593,12 +593,15 @@ submodule (stdlib_linalg) stdlib_linalg_least_squares
         integer(ilp) :: m, n, i
         logical(lk) :: copy_a
         ${rt}$, allocatable :: a_scaled(:,:), b_scaled(:)
-        real(${rk}$) :: sqrt_w(size(w))
+        real(${rk}$), allocatable :: sqrt_w(:)
         character(*), parameter :: this = 'weighted_lstsq'
 
         m = size(a, 1, kind=ilp)
         n = size(a, 2, kind=ilp)
 
+        ! Allocate result (even on error, to prevent segfault on return)
+        allocate(x(n))
+
         ! Validate inputs
         if (size(w, kind=ilp) /= m) then
             err0 = linalg_state_type(this, LINALG_VALUE_ERROR, &
@@ -608,7 +611,7 @@ submodule (stdlib_linalg) stdlib_linalg_least_squares
             return
         end if
 
-        if (any(w <= 0)) then
+        if (any(w <= 0.0_${rk}$)) then
             err0 = linalg_state_type(this, LINALG_VALUE_ERROR, 'Weights must be positive')
             call linalg_error_handling(err0, err)
             if (present(rank)) rank = 0
@@ -623,7 +626,8 @@ submodule (stdlib_linalg) stdlib_linalg_least_squares
         end if
 
         ! Compute sqrt of weights
-        sqrt_w = sqrt(real(w, ${rk}$))
+        allocate(sqrt_w(m))
+        sqrt_w = sqrt(w)
 
         ! Scale A and b
         if (copy_a) then
@@ -678,8 +682,8 @@ submodule (stdlib_linalg) stdlib_linalg_least_squares
         integer(ilp) :: m, n, p, lda, ldb, lwork, info
         logical(lk) :: copy_a, is_prefactored
         ${rt}$, pointer :: amat(:,:)
-        ${rt}$, allocatable :: amat_alloc(:,:), lmat(:,:)
-        ${rt}$, allocatable :: d(:), y(:), work(:)
+        ${rt}$, allocatable, target :: amat_alloc(:,:)
+        ${rt}$, allocatable :: lmat(:,:), d(:), y(:), work(:)
         character(*), parameter :: this = 'generalized_lstsq'
 
         m = size(a, 1, kind=ilp)

test/linalg/test_linalg_lstsq.fypp

@@ -156,7 +156,8 @@ module test_linalg_least_squares
         type(linalg_state_type) :: state
         integer(ilp), parameter :: m = 4, n = 2
         real(${rk}$), parameter :: tol = 100*sqrt(epsilon(0.0_${rk}$))
-        ${rt}$ :: A(m,n), b(m), w(m)
+        ${rt}$ :: A(m,n), b(m)
+        real(${rk}$) :: w(m)
         ${rt}$, allocatable :: x(:)
 
         ! Simple test case
@@ -186,7 +187,8 @@ module test_linalg_least_squares
         type(linalg_state_type) :: state
         integer(ilp), parameter :: m = 4, n = 2
         real(${rk}$), parameter :: tol = 100*sqrt(epsilon(0.0_${rk}$))
-        ${rt}$ :: A(m,n), b(m), w_uniform(m), w_nonuniform(m)
+        ${rt}$ :: A(m,n), b(m)
+        real(${rk}$) :: w_uniform(m), w_nonuniform(m)
         ${rt}$, allocatable :: x_uniform(:), x_weighted(:)
 
         ! Setup problem
@@ -217,7 +219,8 @@ module test_linalg_least_squares
         type(error_type), allocatable, intent(out) :: error
 
         type(linalg_state_type) :: state
-        ${rt}$ :: A(3,2), b(3), w(3)
+        ${rt}$ :: A(3,2), b(3)
+        real(${rk}$) :: w(3)
         ${rt}$, allocatable :: x(:)
 
         A = 1.0_${rk}$