Pinv not Type Stable?

Why is LinearAlgebra.pinv not type stable? Can it be implemented in a type stable way?

Example (thanks to @albheim)

using LinearAlgebra
using JET

function test(M)
    A = rand(M, M + 2)

report = @report_opt test(10)

││┌ pinv(A::Matrix{Float64}; atol::Float64, rtol::Float64) @ LinearAlgebra /Users/julia/.julia/scratchspaces/a66863c6-20e8-4ff4-8a62-49f30b1f605e/agent-cache/default-honeycrisp-R17H3W25T9.0/build/default-honeycrisp-R17H3W25T9-0/julialang/julia-release-1-dot-9/usr/share/julia/stdlib/v1.9/LinearAlgebra/src/dense.jl:1490
│││ runtime dispatch detected: Base.materialize!(%405::Any, %408::Base.Broadcast.Broadcasted{_A, Nothing, typeof(pinv)} where _A<:Union{Nothing, Base.Broadcast.BroadcastStyle})::Any

Do you have an example where it is not type stable?
My test with a normal floating point matrix seems to be stable according to @code_warntype, and I’m not doubting you, it just makes it much easier to investigate with a concrete example since then we don’t have to hunt for what specific type creates an instability.

Amended my question. Thanks for taking a look. Could it be a false positive by JET?

This is what you mean with not being “type stable” I guess. This is what’s reported by @code_warntype

julia> function test(M)
         A = rand(M, M+2)
test (generic function with 1 method)

julia> @code_warntype test(10)
MethodInstance for test(::Int64)
  from test(M) @ Main REPL[10]:1
1 ─ %1 = (M + 2)::Int64
│        (A = Main.rand(M, %1))
│   %3 = Main.pinv(A)::Matrix{Float64}
└──      return %3

EDIT: the thought below turned out to be wrong :wink:

The problem with the runtime dispatch is the following: you are passing an integer which is then used to generate a matrix with dimensions based on the integer’s value. This way, the compiler cannot know in advance the method call for pinv. In other words: the method call and even return value depend on the actual numerical value of M, therefore the method call will be decided during runtime (called runtime dispatch).

Why not? What’s going on inside pinv to cause this, seeing as size is not part of the type?

1 Like

I think it’s rand(), which will create an M x M+2 matrix, so based on the value of M, the return type will be different and that is passed to pinv. But maybe I have a misunderstanding :wink:

I belive you do. Rand is always of type ‘Matrix{Float64}’. I also tested just rand, and that seemed type stable. The issue also crept up in other code, that decidedly didn’t use rand.

Adding to that the size of a standard Julia Array is not part of its type. Only it’s dimensionality is, which is 2 for Matrix.

1 Like

Yes I guess you are right, I was on the wrong path :wink: (let me add that above to avoid confusion). After so many years of Julia, I still get confused by runtime dispatch stuff :laughing:

1 Like

Methods with keyword arguments are implemented to call hidden helper methods that have the written body, so @code_warntype is really difficult to use with them. Well, I dug out the underlying call and am including the part where things get a little unstable, even though the internal arrays and the return value are inferred concretely.

julia> @code_warntype LinearAlgebra.:(var"#pinv#33")(0.0, (eps(real(float(1.1)))*min(3, 5))*iszero(0.0), pinv,rand(3, 5))
MethodInstance for LinearAlgebra.var"#pinv#33"(::Float64, ::Float64, ::typeof(pinv), ::Matrix{Float64})
  from var"#pinv#33"(atol::Real, rtol::Real, ::typeof(pinv), A::AbstractMatrix{T}) where T in LinearAlgebra at /nix/store/4dkh426qldnal1n4m2czlhrmi4qf9ri5-julia-bin-1.8.3/share/julia/stdlib/v1.8/LinearAlgebra/src/dense.jl:1445
│    %40 = (rtol * maxabsA)::Float64
│    %41 = LinearAlgebra.max(%40, atol)::Float64
│          Core.setfield!(tol@_14, :contents, %41)
│    %55 = Base.getproperty(SVD, :S)::Vector{Float64}
│    %56 = LinearAlgebra.maximum(%55)::Float64
│    %57 = (rtol * %56)::Float64
│    %58 = LinearAlgebra.max(%57, atol)::Float64
│          Core.setfield!(tol@_14, :contents, %58)
│    %68 = Core.isdefined(tol@_14, :contents)::Bool
└───       goto #9 if not %68
8 ──       goto #10
9 ──       Core.NewvarNode(:(tol@_21))
└───       tol@_21
10 ┄ %73 = Core.getfield(tol@_14, :contents)::Any
│    %74 = Base.broadcasted(LinearAlgebra.:>, %67, %73)::Any
│          (index = Base.materialize(%74))
│    %76 = Base.dotview(Sinv, index)::Any
│    %77 = Base.getproperty(SVD, :S)::Vector{Float64}
│    %78 = LinearAlgebra.view(%77, index)::Any
│    %79 = Base.broadcasted(LinearAlgebra.pinv, %78)::Base.Broadcast.Broadcasted{_A, Nothing, typeof(pinv)} where _A<:Union{Nothing, Base.Broadcast.BroadcastStyle}
│          Base.materialize!(%76, %79)

The corresponding source code section is:

        tol = max(rtol * maxabsA, atol)
    tol         = max(rtol*maximum(SVD.S), atol)
    index       = SVD.S .> tol
    Sinv[index] .= pinv.(view(SVD.S, index))

It seems as if tol is an uninferred boxed value (Core.Box) and it’s not even certain if it is defined (%68: tol@_14 vs tol@_21). I don’t really know why because in both lines (%41, %58) the final max call is inferred concretely ::Float64 before being stored in the box.


pinv reassigns tol after it is captured by a closure, I think that’s the reason. Closures often ruin inferrability because it’s understandably hard to infer a variable shared by 2 separate functions that get compiled separately. I don’t know all the cases where it is inferred (guaranteed reassignment before closure instantiation) and where it is not (conditional reassignment before closure, any reassignment after the closure).

I @evaled a pinv2 into LinearAlgebra that replaced the closure with a function that doesn’t capture tol: B[indB] .= ((x, toll) -> abs(x) > toll ? pinv(x) : zero(x)).(dA, tol). I haven’t compared the performance of these anonymous functions in isolation, so I don’t think it’s optimal. But the box vanished:

julia> @code_warntype LinearAlgebra.:(var"#pinv2#251")(0.0, (eps(real(float(1.1)))*min(3, 5))*iszero(0.0), LinearAlgebra.pinv2,rand(3, 5))
MethodInstance for LinearAlgebra.var"#pinv2#251"(::Float64, ::Float64, ::typeof(LinearAlgebra.pinv2), ::Matrix{Float64})
  from var"#pinv2#251"(atol::Real, rtol::Real, ::typeof(LinearAlgebra.pinv2), A::AbstractMatrix{T}) where T in LinearAlgebra at REPL[16]:1

EDIT: the least disruptive change to accomplish the same inferability would be to keep the closure, but alter the captured variable so it’s not reassigned afterward.

        tol_diag_branch = max(rtol * maxabsA, atol)
        B[indB] .= (x -> abs(x) > tol_diag_branch ? pinv(x) : zero(x)).(dA)
        # do not reassign tol_diag_branch, so parser lets compiler infer it

Does this have performance implications? I assume not really since it’s doing only one dispatch per call.

Not really; computing the pinv — even for very small arrays — is orders of magnitude more work than the cost of this internal instability. Where it might have an impact is in static compilation and potentially being a source for triggering invalidations, so it’s still a good thing to address. @Benny you should post that change as a pull request! :slight_smile:

The TL;DR of why this is unstable is:

  • tol is assigned to twice in a slightly complicated manner — once in an if branch and once outside.
  • tol is also captured inside a closure, which theoretically has to support getting an updated value from the second assignment. In this case it’s immediately evaluated (and is in a completely independent code path), but the front-end isn’t grokking that.

So this hits the infamous issue #15276. An even simpler fix is to use two completely separate names for tol for the two different codepaths.


Thanks, I learned quite a lot in this thread. I’ll run my own implementation for now since I don’t expect myself to hand me any diagonal matrices.

1 Like

Yh it would be useful to get this fixed as it seems to cause pinv to not work on the GPU, Support for LinearAlgebra.pinv · Issue #2070 · JuliaGPU/CUDA.jl · GitHub.


So I ran some benchmarks with the changes @mbauman suggested, and types do indeed get inferred correctly now. The two implementations are indistinguishable from a performance perspective.

"Type Stable Pinv"

BenchmarkTools.Trial: 4059 samples with 1 evaluation.
 Range (min … max):  1.138 ms …   6.104 ms  ┊ GC (min … max): 0.00% … 79.07%
 Time  (median):     1.177 ms               ┊ GC (median):    0.00%
 Time  (mean ± σ):   1.231 ms ± 356.126 μs  ┊ GC (mean ± σ):  3.26% ±  8.18%

  █▆▁                                                         ▁
  ██████▇▇▄▅▅▅▅▄▄▄▄▃▄▃▁▁▄▁▁▃▁▁▁▁▁▁▁▁▁▃▃▁▁▁▁▁▃▁▁▁▁▁▁▁▁▁▁▃▅▁▅▃▅ █
  1.14 ms      Histogram: log(frequency) by time      3.71 ms <

 Memory estimate: 645.53 KiB, allocs estimate: 19.

BenchmarkTools.Trial: 4008 samples with 1 evaluation.
 Range (min … max):  1.143 ms …   7.397 ms  ┊ GC (min … max): 0.00% … 69.65%
 Time  (median):     1.178 ms               ┊ GC (median):    0.00%
 Time  (mean ± σ):   1.246 ms ± 379.645 μs  ┊ GC (mean ± σ):  3.12% ±  8.04%

  ███▇█▇▇▆▇▇▆▅▆▄▅▃▅▄▄▁▃▄▃▄▄▄▄▄▄▄▃▃▁▁▃▃▃▃▁▃▃▁▃▃▁▁▄▁▃▁▁▄▃▄▄▅▄▄▅ █
  1.14 ms      Histogram: log(frequency) by time      3.67 ms <

 Memory estimate: 645.81 KiB, allocs estimate: 27.

The fix has been merged into master, Remove boxing in pinv by Zentrik · Pull Request #51351 · JuliaLang/julia · GitHub.