Eigenvalues and eigenvectors #1334
Conversation
|
thanks @awni my question is:
|
python/src/linalg.cpp
Outdated
| nb::sig( | ||
| "def eigvalsh(a: array, upper: bool = True, *, stream: Union[None, Stream, Device] = None) -> array"), | ||
| R"pbdoc( | ||
| Compute the eigenvalues of a complex Hermitian or real symmetric matrix. |
There was a problem hiding this comment.
We don't do a brilliant job of this in the other linalg functions, but we should probably check if complex inputs work and raise a not implemented error if not.
mlx/primitives.h
Outdated
| @@ -2158,4 +2158,23 @@ class Cholesky : public UnaryPrimitive { | |||
| bool upper_; | |||
| }; | |||
|
|
|||
| class Eigvalsh : public UnaryPrimitive { | |||
| public: | |||
| explicit Eigvalsh(Stream stream, bool upper, bool compute_vectors) | |||
There was a problem hiding this comment.
Given it can compute eigenvectors, this can probably be called Eigh rather than Eigvalsh?
| // Delegate to the eigenvalue decomposition taking into account differences in | ||
| // LAPACK implementations (basically how to pass the 'jobz' and 'uplo' strings | ||
| // to fortran). | ||
| int ssyevd_wrapper(char jobz, char uplo, float* matrix, float* w, int N) { |
There was a problem hiding this comment.
I don't think it needs to be in this PR, but it would be nice to add support for arbitrary matrices (not just hermitian/symmetric ones).
Hopefully that could just be a flag on the Eig primitive that then uses a different lapack incantation?
mlx/linalg.h
Outdated
| @@ -74,4 +74,8 @@ array pinv(const array& a, StreamOrDevice s = {}); | |||
|
|
|||
| array cholesky_inv(const array& a, bool upper = false, StreamOrDevice s = {}); | |||
|
|
|||
| array eigvalsh(const array& a, bool upper = false, StreamOrDevice s = {}); | |||
|
|
|||
| array eigh(const array& a, bool upper = false, StreamOrDevice s = {}); | |||
There was a problem hiding this comment.
Shouldn't that return both the eigenvalues and eigenvectors like numpy?
mlx/primitives.h
Outdated
| @@ -2158,4 +2158,36 @@ class Cholesky : public UnaryPrimitive { | |||
| bool upper_; | |||
| }; | |||
|
|
|||
| class Eigvalsh : public UnaryPrimitive { | |||
There was a problem hiding this comment.
Sorry if @barronalex already made this suggestion: but I think it makes sense to merge the Eighvalsh and Eigh into a single primitive. And have the ops use the same primitive but just return only the eigenvalues in the case of eighvalsh.
It looks like the work is done anyway.. and the underlying implementations are basically identical.
|
@awni is the current implementation more like what you mean? |
| @@ -81,6 +81,7 @@ DEFAULT_MULTI(SVD) | |||
| DEFAULT(Transpose) | |||
| DEFAULT(Inverse) | |||
| DEFAULT(Cholesky) | |||
| DEFAULT_MULTI(EighPrimitive) | |||
There was a problem hiding this comment.
Please rename this to Eigh for consistency with other primitive names.
|
Yes I think it's better this way. I'm debating if we should bother with the |
Proposed changes
Please include a description of the problem or feature this PR is addressing. If there is a corresponding issue, include the issue #.
Checklist
Put an
xin the boxes that apply.pre-commit run --all-filesto format my code / installed pre-commit prior to committing changes