Create blocked Jacobi method for eigen decomposition #1510
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Draft commit to introduce the idea. Todo: * Handle complex numbers * Reject malformed matrices
polvalente
reviewed
Jun 12, 2024
nx/lib/eigh_block.ex
Outdated
Comment on lines
59
to
68
| defn eigh(matrix) do | ||
| matrix | ||
| |> Nx.revectorize([collapsed_axes: :auto], | ||
| target_shape: {Nx.axis_size(matrix, -2), Nx.axis_size(matrix, -1)} | ||
| ) | ||
| |> decompose() | ||
| |> then(fn {w, v} -> | ||
| revectorize_result({w, v}, matrix) | ||
| end) | ||
| end |
Contributor
There was a problem hiding this comment.
I believe this should be the only defn in this module and the others would be defnp. Or something close to that.
polvalente
reviewed
Jun 12, 2024
nx/lib/eigh_block.ex
Outdated
| end | ||
|
|
||
| # Initialze tensors to hold eigenvectors | ||
| v_tl = Nx.eye(mid, type: :f32) |
Contributor
There was a problem hiding this comment.
why force f32 here? Is this another case where the algorithm just fails on f64?
Perhaps this should be masked underneath the implementation if it's the case.
polvalente
reviewed
Jun 12, 2024
nx/lib/eigh_block.ex
Outdated
| # | ||
| # The inner loop performs "sweep" rounds of n - 1, which is enough permutations to allow | ||
| # all sub matrices to share the needed values. | ||
| {_, _, tl, _tr, _bl, br, v_tl, v_tr, v_bl, v_br, _} = |
Contributor
There was a problem hiding this comment.
You can use a pattern for organizing the while state that we do quite a lot:
{{tl, br, v_tl, v_tr, v_bl, v_br}, _} where you leave the outputs in a first-position tuple, and the other state in a second position, so pattern matching on the statement is easier, as well as understanding what's output and what's not
polvalente
approved these changes
Jun 12, 2024
Contributor
polvalente
left a comment
polvalente
left a comment
There was a problem hiding this comment.
Looking good. I'm also leaving a few stylistic suggestions for readability
nx/lib/eigh_block.ex
Outdated
| t = Nx.sqrt(1 + Nx.pow(tau, 2)) | ||
| t = Nx.select(Nx.greater_equal(tau, 0), 1 / (tau + t), 1 / (tau - t)) | ||
|
|
||
| pred = Nx.less_equal(Nx.abs(b), 0.1 * 1.0e-4 * Nx.min(Nx.abs(a), Nx.abs(c))) |
Contributor
There was a problem hiding this comment.
Suggested change
| pred = Nx.less_equal(Nx.abs(b), 0.1 * 1.0e-4 * Nx.min(Nx.abs(a), Nx.abs(c))) | |
| pred = Nx.abs(b) <= 1.0e-5 * Nx.min(Nx.abs(a), Nx.abs(c)) |
Co-authored-by: Paulo Valente <16843419+polvalente@users.noreply.github.com>
Co-authored-by: Paulo Valente <16843419+polvalente@users.noreply.github.com>
…/nx into cgreen/block-jacobi-eigh
Contributor
|
Merging this now! For reference, these are the benchmarks on my desktop, where CudaHost |
polvalente
changed the title
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
#1027
The current implementation of
eighusing a 10x10 symmetric matrix takes about450msfor a 20x20 matrix and154sfor a 100x100 matrix, while the new implementation takes0.5msand11msrespectively.This is a
defnversion of the method used by XLA: https://github.com/openxla/xla/blob/main/xla/service/eigh_expander.ccThere is still a todo list and code cleanup/drying to do, but I wanted to pitch this before getting to far into the process. While this method has a static submatrix size with no recursion, this approach can be built on to recreate the recursive blocked-eigh used by JAX. This approach had less complexity and seemed like a nice way to make
eighperformant without having to exactly copy the JAX method.The gist of the method is to break the matrix into four submatrices and apply the jacobi rotations across all rows and cols each iteration and then joining the results.
Draft commit to introduce the idea.
Todo:
Current issues:
Please let me know if this is of any use! <3