zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-11 22:44:38 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

CPU mx.linalg.cholesky_inverse and mx.linalg.tri_inv (#1307)

Browse Source 
* add cholesky inv + tri inv

* always run tri_inv on cpu

* consistent naming
This commit is contained in:
Alex Barron
2024-08-08 15:18:02 -07:00
committed by GitHub
parent 780c197f95
commit 32668a7317
7 changed files with 267 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files

65
python/src/linalg.cpp
View File

@@ -261,6 +261,31 @@ void init_linalg(nb::module_& parent_module) {
array: ``ainv`` such that ``dot(a, ainv) = dot(ainv, a) = eye(a.shape[0])``
)pbdoc");
m.def(
"tri_inv",
&tri_inv,
"a"_a,
"upper"_a,
nb::kw_only(),
"stream"_a = nb::none(),
nb::sig(
"def tri_inv(a: array, upper: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"),
R"pbdoc(
Compute the inverse of a triangular square matrix.
This function supports arrays with at least 2 dimensions. When the input
has more than two dimensions, the inverse is computed for each matrix
in the last two dimensions of ``a``.
Args:
a (array): Input array.
upper (array): Whether the array is upper or lower triangular. Defaults to ``False``.
stream (Stream, optional): Stream or device. Defaults to ``None``
in which case the default stream of the default device is used.
Returns:
array: ``ainv`` such that ``dot(a, ainv) = dot(ainv, a) = eye(a.shape[0])``
)pbdoc");
m.def(
"cholesky",
&cholesky,
"a"_a,
@@ -286,8 +311,46 @@ void init_linalg(nb::module_& parent_module) {
in which case the default stream of the default device is used.
Returns:
array: If ``upper = False``, it returns a lower trinagular ``L`` matrix such
array: If ``upper = False``, it returns a lower triangular ``L`` matrix such
that ``dot(L, L.T) = a``. If ``upper = True``, it returns an upper triangular
``U`` matrix such that ``dot(U.T, U) = a``.
)pbdoc");
m.def(
"cholesky_inv",
&cholesky_inv,
"a"_a,
"upper"_a = false,
nb::kw_only(),
"stream"_a = nb::none(),
nb::sig(
"def cholesky_inv(L: array, upper: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"),
R"pbdoc(
Compute the inverse of a real symmetric positive semi-definite matrix using it's Cholesky decomposition L.
Let A be a real symmetric positive semi-definite matrix and L its Cholesky definition such that:
.. math::
\begin{aligned}
\mathbf{A} = \mathbf{L}\mathbf{L}^T
\end{aligned}
This function computes :math:`\mathbf{A}^{-1}`.
This function supports arrays with at least 2 dimensions. When the input
has more than two dimensions, the Cholesky inverse is computed for each matrix
in the last two dimensions of ``L``.
If the input matrix is not a triangular matrix behaviour is undefined.
Args:
L (array): Input array.
upper (bool, optional): If ``True``, return the upper triangular Cholesky factor.
If ``False``, return the lower triangular Cholesky factor. Default: ``False``.
stream (Stream, optional): Stream or device. Defaults to ``None``
in which case the default stream of the default device is used.
Returns:
array: :math:`A^{-1}` where :math:`\mathbf{A} = \mathbf{L}\mathbf{L}^T`.
)pbdoc");
}