Module: tfp.experimental.substrates.jax.math.linalg

Functions for common linear algebra operations.


cholesky_concat(...): Concatenates chol @ chol.T with additional rows and columns.

cholesky_update(...): Returns cholesky of chol @ chol.T + multiplier * u @ u.T.

fill_triangular(...): Creates a (batch of) triangular matrix from a vector of inputs.

fill_triangular_inverse(...): Creates a vector from a (batch of) triangular matrix.

lu_matrix_inverse(...): Computes a matrix inverse given the matrix's LU decomposition.

lu_reconstruct(...): The inverse LU decomposition, X == lu_reconstruct(*

lu_reconstruct_assertions(...): Returns list of assertions related to lu_reconstruct assumptions.

lu_solve(...): Solves systems of linear eqns A X = RHS, given LU factorizations.

pivoted_cholesky(...): Computes the (partial) pivoted cholesky decomposition of matrix.

sparse_or_dense_matmul(...): Returns (batched) matmul of a SparseTensor (or Tensor) with a Tensor.

sparse_or_dense_matvecmul(...): Returns (batched) matmul of a (sparse) matrix with a column vector.