-
Notifications
You must be signed in to change notification settings - Fork 950
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Eigenvalues and eigenvectors #1334
base: main
Are you sure you want to change the base?
Conversation
thanks @awni my question is:
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This is nicely done and a great addition -- thanks!
Left a few comments/suggestions.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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
x
in the boxes that apply.pre-commit run --all-files
to format my code / installed pre-commit prior to committing changes