Givens orthogonal layer by VolodyaCO · Pull Request #57 · dwavesystems/dwave-pytorch-plugin

VolodyaCO · 2025-12-19T00:01:02Z

This PR adds an orthogonal layer given by Givens rotations, using the parallel algorithm described by Firas in https://arxiv.org/abs/2106.00003, which gives a forward complexity of O(n) and backward complexity of O(n log(n)), even though there are O(n^2) rotations.

This PR still is in draft. I wrote it for even n. Probably some more unit tests are to be done, but I am quite lazy (will do it after all math is checked for odd n).

VolodyaCO · 2025-12-19T00:25:33Z

I somehow broke @kevinchern's tests, what the hell...

kevinchern · 2025-12-19T00:39:51Z

I somehow broke @kevinchern's tests, what the hell...

@VolodyaCO which tests? I'm seeing test_forward_agreement and test_backward_agreement failures on this CI test

VolodyaCO · 2025-12-19T14:34:49Z

I somehow broke @kevinchern's tests, what the hell...

@VolodyaCO which tests? I'm seeing test_forward_agreement and test_backward_agreement failures on this CI test

I forgot to update my tests to float64 precision. Now that I've done it, it's weird that all of the current failing tests are failing on

  File "/Users/distiller/project/tests/test_nn.py", line 144, in test_LinearBlock
    self.assertTrue(model_probably_good(model, (din,), (dout,)))

kevinchern · 2025-12-19T15:41:10Z

I somehow broke @kevinchern's tests, what the hell...

@VolodyaCO which tests? I'm seeing test_forward_agreement and test_backward_agreement failures on this CI test

I forgot to update my tests to float64 precision. Now that I've done it, it's weird that all of the current failing tests are failing on
  File "/Users/distiller/project/tests/test_nn.py", line 144, in test_LinearBlock
    self.assertTrue(model_probably_good(model, (din,), (dout,)))

Ahhhhhh. OK Theo also flagged this at #50 . It's a poorly-written test.. you can ignore it.

thisac · 2025-12-30T00:50:29Z

+        angles, blocks, Ufwd_saved = ctx.saved_tensors
+        Ufwd = Ufwd_saved.clone()
+        M = grad_output.t()  # dL/dU, i.e., grad_output is of shape (n, n)
+        n = M.size(1)
+        block_size = n // 2
+        A = torch.zeros((block_size, n), device=angles.device, dtype=angles.dtype)


Same here re lowercase for Ufwd, M, and A. Avoids incorrect colour highlighting in themes.

Hmmm, I didn't read this about the incorrect colour highlighting before I made my previous comment. I still think that it is easier to read the algorithm alongside the code if the use of lower/upper case match. For example, lower case m is usually used for an integer variable, not a tensor.

VolodyaCO · 2025-12-31T19:16:03Z

After a bit of git wrangling, I was able to clean my whole mess of merge commits 😆.

anahitamansouri

This is a nice PR Vlad. It took me a while to go over the paper and this PR :) The only thing is the tests that are failing. Thanks for the great work.

kevinchern

Did a quick pass to provide some feedback before taking some time to take a deep dive into the paper.

kevinchern · 2026-01-08T19:46:05Z

+    for _ in range(n - 1):
+        pairs = circle_method(sequence)
+        if is_odd:
+            # Remove pairs involving the dummy dimension:


Suggested change

# Remove pairs involving the dummy dimension:

# Remove pairs involving the dummy dimension

I was gonna ask why remove the colon?

kevinchern · 2026-04-17T18:05:01Z

@VolodyaCO can you rebase on main?

VolodyaCO

@kevinchern I have addressed all comments. Should be ready to merge now.

thisac · 2026-05-05T23:04:05Z

+__all__ = ["GivensRotation"]
+
+
+class _RoundRobinGivens(torch.autograd.Function):


The documentation is quite thorough for a hidden class. Would this perhaps make sense moving to a nn.functions namespace and removing the underscore? @kevinchern

@kevinchern bumping this

I don't think it will be used beyond this class, so it's probably safe to keep her. If there's demand for it as a public function, we can move it out.

thisac · 2026-05-05T23:08:44Z

+        # Initialize U^fwd from forward pass output U. Mathematically, U^fwd represents U^{1:k-1}
+        # at block k, defined in equation (11). It is post-multiplied by G_bk^T at each block
+        # iteration to "remove the effect of the block's rotations" (Section 4.1, paragraph before
+        # equation (12)). This corresponds to the update from U^{1:k} back to U^{1:k-1}.


I haven't looked at the reference link, but these comments seem a bit messy (for lack of a better term). If the algorithm in the link describes the process in a clear way, I'd perhaps shorten these comments a bit, although references to equations are always good.

Also, separation by empty lines always makes things look a bit tidier.

I have added separation by empty lines to make things look tidier. I added these inline comments at @kevinchern 's request. I do agree that it is enough for the interested reader to read the paper. I also think it's helpful to have the inline comments too.

shorten these comments a bit

I second this opinion

kevin's request

To clarify my suggestion: it will be very helpful for future maintenance to include references to "where" these equations appear. The idea is for the developer to be able to easily reference the material and not necessarily understand the method from comments.

I see your point. Comments have now been shortened.

kevinchern

Unit tests hinge on NaiveGivensRotationLayer being correct but there aren't correctness tests for it. It checks for serial versus parallel implementations.
We need a unit test to check the forward and backward methods are correct, e.g., work out the expected forward/backward values a small input (n=3?)

kevinchern · 2026-05-26T00:44:22Z

+        # Initialize U^fwd from forward pass output U. Mathematically, U^fwd represents U^{1:k-1}
+        # at block k, defined in equation (11). It is post-multiplied by G_bk^T at each block
+        # iteration to "remove the effect of the block's rotations" (Section 4.1, paragraph before
+        # equation (12)). This corresponds to the update from U^{1:k} back to U^{1:k-1}.


shorten these comments a bit

I second this opinion

kevin's request

To clarify my suggestion: it will be very helpful for future maintenance to include references to "where" these equations appear. The idea is for the developer to be able to easily reference the material and not necessarily understand the method from comments.

Address orthogonal module PR feedback.

kevinchern

LGTM; comments have been addressed.

VolodyaCO requested a review from kevinchern December 19, 2025 00:01

VolodyaCO self-assigned this Dec 19, 2025

VolodyaCO added the enhancement New feature or request label Dec 19, 2025

VolodyaCO requested review from anahitamansouri and thisac December 19, 2025 00:25

kevinchern reviewed Dec 19, 2025

View reviewed changes

Comment thread tests/helper_models.py Outdated

kevinchern reviewed Dec 19, 2025

View reviewed changes

Comment thread tests/test_nn.py

VolodyaCO marked this pull request as ready for review December 26, 2025 17:28

thisac reviewed Dec 30, 2025

View reviewed changes

VolodyaCO force-pushed the givens-rotation branch from f18b476 to 46062e0 Compare December 31, 2025 19:12

anahitamansouri reviewed Jan 5, 2026

View reviewed changes

Comment thread dwave/plugins/torch/nn/modules/orthogonal.py Outdated

VolodyaCO force-pushed the givens-rotation branch from 46062e0 to 7f76571 Compare January 5, 2026 18:38

kevinchern requested changes Jan 8, 2026

View reviewed changes

VolodyaCO force-pushed the givens-rotation branch 3 times, most recently from 1146ebb to df400fb Compare January 14, 2026 14:41

VolodyaCO force-pushed the main branch from c207935 to eb65b87 Compare March 5, 2026 23:21

VolodyaCO force-pushed the givens-rotation branch from df400fb to 8a00932 Compare April 20, 2026 15:13

kevinchern requested changes Apr 21, 2026

View reviewed changes

Comment thread dwave/plugins/torch/nn/modules/orthogonal.py

Comment thread tests/helper_models.py Outdated

Comment thread tests/test_nn.py

VolodyaCO force-pushed the givens-rotation branch 3 times, most recently from 6587b89 to 97c2b8e Compare April 27, 2026 21:33

kevinchern closed this Apr 27, 2026

kevinchern reopened this Apr 27, 2026

VolodyaCO closed this Apr 27, 2026

VolodyaCO reopened this Apr 27, 2026

VolodyaCO force-pushed the givens-rotation branch from 97c2b8e to 36a15c8 Compare April 27, 2026 22:25

VolodyaCO force-pushed the givens-rotation branch from 36a15c8 to 016cc66 Compare May 4, 2026 22:38

VolodyaCO commented May 4, 2026

View reviewed changes

Comment thread tests/test_nn.py

Comment thread dwave/plugins/torch/nn/modules/orthogonal.py

thisac reviewed May 5, 2026

View reviewed changes

VolodyaCO force-pushed the givens-rotation branch from 016cc66 to 9f31f15 Compare May 7, 2026 15:38

kevinchern requested changes May 26, 2026

View reviewed changes

Add Givens orthogonal layer

c0cd080

Address orthogonal module PR feedback.

VolodyaCO force-pushed the givens-rotation branch from 9f31f15 to c0cd080 Compare June 29, 2026 18:31

kevinchern approved these changes Jul 13, 2026

View reviewed changes

	# Remove pairs involving the dummy dimension:
	# Remove pairs involving the dummy dimension

		__all__ = ["GivensRotation"]


		class _RoundRobinGivens(torch.autograd.Function):

Uh oh!

Conversation

VolodyaCO commented Dec 19, 2025

Uh oh!

VolodyaCO commented Dec 19, 2025

Uh oh!

Uh oh!

Uh oh!

kevinchern commented Dec 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

VolodyaCO commented Dec 19, 2025

Uh oh!

kevinchern commented Dec 19, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

VolodyaCO commented Dec 31, 2025

Uh oh!

anahitamansouri left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

kevinchern left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

kevinchern commented Apr 17, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

VolodyaCO left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kevinchern left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

kevinchern commented Dec 19, 2025 •

edited

Loading