Decomposition.hs

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ExplicitForAll #-}
{-# LANGUAGE TypeApplications #-}

-- | User-facing EDSL API for dense linear algebra decompositions.
module EigenHHLO.EDSL.Decomposition
    ( svd
    , qr
    , q
    , eig
    , chol
    , lu
    ) where

import GHC.TypeLits (KnownNat)
import HHLO.Core.Types (DType(..))
import HHLO.IR.Builder (Builder, Tensor)

import EigenHHLO.Core.Types (EigenSession)
import qualified EigenHHLO.IR.Cholesky as IR
import qualified EigenHHLO.IR.Eigenvalue as IR
import qualified EigenHHLO.IR.LU as IR
import qualified EigenHHLO.IR.QR as IR
import qualified EigenHHLO.IR.SVD as IR

-- | Singular value decomposition.
svd :: forall m n k. (KnownNat m, KnownNat n, KnownNat k)
    => EigenSession
    -> Tensor '[m, n] 'F64
    -> Builder (Tensor '[m, k] 'F64, Tensor '[k] 'F64, Tensor '[k, n] 'F64)
svd = IR.svdBuilder

-- | QR factorization: A = Q · R.
-- Returns (A overwritten with reflectors, tau).
qr :: forall m n k. (KnownNat m, KnownNat n, KnownNat k)
   => EigenSession
   -> Tensor '[m, n] 'F64
   -> Builder (Tensor '[m, n] 'F64, Tensor '[k] 'F64)
qr = IR.qrBuilder

-- | Generate explicit Q from QR reflectors.
q :: forall m n k. (KnownNat m, KnownNat n, KnownNat k)
  => EigenSession
  -> Tensor '[m, n] 'F64
  -> Tensor '[k] 'F64
  -> Builder (Tensor '[m, m] 'F64)
q = IR.qBuilder

-- | Symmetric eigenvalue decomposition.
eig :: forall n. (KnownNat n)
    => EigenSession
    -> Tensor '[n, n] 'F64
    -> Builder (Tensor '[n] 'F64, Tensor '[n, n] 'F64)
eig = IR.eigBuilder

-- | Cholesky factorization.
chol :: forall n. (KnownNat n)
     => EigenSession
     -> Tensor '[n, n] 'F64
     -> Builder (Tensor '[n, n] 'F64)
chol = IR.cholBuilder

-- | LU factorization with partial pivoting.
lu :: forall m n k. (KnownNat m, KnownNat n, KnownNat k)
   => EigenSession
   -> Tensor '[m, n] 'F64
   -> Builder (Tensor '[m, n] 'F64, Tensor '[k] 'I32)
lu = IR.luBuilder