diff --git a/src/CPU.jl b/src/CPU.jl index 4df7e01..892acc5 100644 --- a/src/CPU.jl +++ b/src/CPU.jl @@ -324,7 +324,8 @@ function bidiag_forward!(u::Vector{T}, g::FixedEffectCoefficients, fem::FixedEff # serial: everything in one pass, gathers accumulated directly into g s = _bidiag_chunk!(u, gs, blocks, qrowss, vs, T(c), eachindex(u)) end - return T(sqrt(s)) + # The asserted norm type prevents boxing in LSMR's scalar recurrences. + return T(sqrt(s))::T end # u[i] ← c * u[i] + Σ_blocks fit_i; accumulates each block's gather into diff --git a/test/solve.jl b/test/solve.jl index ac70586..8738bc9 100644 --- a/test/solve.jl +++ b/test/solve.jl @@ -21,6 +21,16 @@ r_ols = [-0.2015993617092453, 0.2015993617092464, -0.2015993617092463, 0.2015 @test iterz == 0 @test all(iszero, rz) +@testset "Inferred bidiagonalization norm" begin + for T in (Float32, Float64) + solver = FixedEffects.AbstractFixedEffectSolver{T}(fes, uweights(T, length(x)), Val{:cpu}) + u = T.(x) + fill!(solver.v, one(T)) + @test (@inferred FixedEffects.bidiag_forward!( + u, solver.g, solver.m, solver.v, zero(T))) isa T + end +end + @testset "maxiter semantics" begin (r0, iter0, conv0) = @test_logs (:warn, r"solve_residuals!") solve_residuals!(deepcopy(x), fes; maxiter = 0) @test iter0 == 0