LinearAlgebra meta-issue

[wsmoses]

Probably good to track how much of LinearAlgebra.jl coverage we have atm. List of Operations (several are missing atm):

• adjoint
• adjoint!
• axpy
• axpby
• copy_adjoint!
• copy_transpose!
• cond
• condskeel
• cross
• det
• logdet
• logabsdet
• diag
• diagm
• dot
• givens
• hermitianpart/hermitianpart!
• isdiag
• ishermitian
• isposdef/isposdef!
• issuccess
• issymmetric
• istril
• istriu
• kron/kron!
• ldiv!
• lmul!
• lowrankdowndate/lowrankdowndate!
• lowrankupdate/lowrankupdate!
• lyap
• mul!
• norm
• normalize/normalize!
• nullspace
• opnorm
• inv
• pinv
• rank
• rdiv!
• reflect!
• rotate!
• stride1
• sylvester
• tr
• tril
• tril!
• triu
• triu!
• transpose/transpose!
• Factorization
  • factorize
  • cholesky/cholesky!
  • lu/lu!
  • lq/lq!
  • qr/qr!
  • eigen/eigen!
  • eigvals/eigvals!
  • eigmax
  • eigmin
  • eigvecs
  • svd/svd!
  • svdvals/svdvals!
  • schur/schur!
  • ordschur/ordschur!
  • ldlt/ldlt!
  • bunchkaufman/bunchkaufman!
  • hessenberg/hessenberg!

[mofeing]

So I've added the rest of functions defined in LinearAlgebra (except for the BLAS and LAPACK submodules).

Some notes:

• I usually put the 2 in-place and "out-place" versions because it might be of interest to implement both. Reasoning is that in some of the out-place methods might just allocate a matrix and call the in-place version, which would force us to create a stablehlo.constant and later pass it. Sometimes this might not be required like with cholesky/cholesky!, because StableHLO already has an op that directly returns the array without modifying in-place.
• I'm not sure about the strategy we should follow for the BLAS/LAPACK implementations: should we call custom_call or should we implement our own methods? e.g. for SVD, should we call LAPACK or have our own numerical routine on top of other HLO ops. IMO we should go for custom_call on BLAS/LAPACK pointing to libblastrampoline routines, but we could replace some for HLO ops like use dot_general instead of blas_*gemm.

@wsmoses i'm really interested into having svd, qr, lu and eigen (and their derivatives) for my research. i'm already working on having the custom_call op, but how can we register the derivatives for them?

[wsmoses]

we should eventualy add the custom derivative interface

[wsmoses]

for the time being I think we can ignore the direct BLAS/LAPACK calls for now, but we definitely should try to flush out the rest of linalg

[avik-pal]

XLA (not in StableHLO) has QR (https://github.com/openxla/xla/blob/492bff6877eca4710c9f008836e93d4aad6a81dd/xla/hlo/builder/lib/qr.cc#L35) and LU (https://github.com/openxla/xla/blob/492bff6877eca4710c9f008836e93d4aad6a81dd/xla/hlo/builder/lib/lu_decomposition.cc) implemented. Can we call them directly somehow?

[mofeing]

And SVD (https://github.com/openxla/xla/blob/492bff6877eca4710c9f008836e93d4aad6a81dd/xla/hlo/builder/lib/svd.h) and Eigendecomposition for self-adjoint matrices (https://github.com/openxla/xla/blob/492bff6877eca4710c9f008836e93d4aad6a81dd/xla/hlo/builder/lib/self_adjoint_eig.h)

I believe they are implemented as async calls so we can pass the output of a compiled function to them, but not sure if we can call them from a stablehlo.custom_call? also, how can we add the diff rules for them?

[wsmoses]

So we should be able to just do stablehlo.custom_call there like https://github.com/jax-ml/jax/blob/3f24dfd234054871ea6682ca2116fc20dd7f36c1/jax/_src/lax/linalg.py#L1805

[avik-pal]

I realized (atleast the comments say so) some of the ones in XLA are specific for TPUs.

[mofeing]

I realized (atleast the comments say so) some of the ones in XLA are specific for TPUs.

Oh no, do you know which ones?

[avik-pal]

https://github.com/openxla/xla/blob/492bff6877eca4710c9f008836e93d4aad6a81dd/xla/hlo/builder/lib/lu_decomposition.cc#L53-L54