zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Custom Metal Kernels from Python (#1325)

Browse Source 
* start

* simple kernels working

* restructure

* inverse example working

* docs + fixes

* missing file

* fix imports

* address comments

* add docs + fix test

* Review comments + refactor to a single function

* update docs

* remove hashing

* fix contig bug in test

* back to a class

* trailing whitespace

* fix tests

* match c++ and python apis

* add link + make args kw_only
This commit is contained in:
Alex Barron
2024-08-22 13:46:29 -07:00
committed by GitHub
parent df3233454d
commit 0fd2a1f4b0
12 changed files with 793 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
python/src/linalg.cpp
View File

@@ -325,9 +325,9 @@ void init_linalg(nb::module_& parent_module) {
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.
Compute the inverse of a real symmetric positive semi-definite matrix using it's Cholesky decomposition.
Let A be a real symmetric positive semi-definite matrix and L its Cholesky definition such that:
Let :math:`\mathbf{A}` be a real symmetric positive semi-definite matrix and :math:`\mathbf{L}` its Cholesky decomposition such that:
.. math::
@@ -339,7 +339,7 @@ void init_linalg(nb::module_& parent_module) {
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``.
in the last two dimensions of :math:`\mathbf{L}`.
If the input matrix is not a triangular matrix behaviour is undefined.
@@ -351,6 +351,6 @@ void init_linalg(nb::module_& parent_module) {
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`.
array: :math:`\mathbf{A^{-1}}` where :math:`\mathbf{A} = \mathbf{L}\mathbf{L}^T`.
)pbdoc");
}