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:36:26Z

    def test_store_config(self):
        with self.subTest("Simple case"):
+
            class MyModel(torch.nn.Module):


Remove formatting changes. Is this "black" formatting?

Yes. I have it by default on my vscode

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:41:23Z

+    Returns:
+        list[list[tuple[int, int]]]: Blocks of edges for parallel Givens rotations.
+
+    Note:


Better as a note directive: https://www.sphinx-doc.org/en/master/usage/restructuredtext/directives.html#directive-note

Where should I put this? in the release notes? or in the docstring itself?

Simply change the Note: to

.. note:: Lorem ipsum...

which would render a note box if we generate docs with Sphinx.

I have done this now

thisac · 2025-12-30T00:44:33Z

+            angles (torch.Tensor): A ((n - 1) * n // 2,) shaped tensor containing all rotations
+                between pairs of dimensions.
+            blocks (torch.Tensor): A (n-1, n//2, 2) shaped tensor containing the indices that
+                specify rotations between pairs of dimensions. Each of the n-1 blocks contains n//2
+                pairs of independent rotations.


Code formatting?

Suggested change

angles (torch.Tensor): A ((n - 1) * n // 2,) shaped tensor containing all rotations

between pairs of dimensions.

blocks (torch.Tensor): A (n-1, n//2, 2) shaped tensor containing the indices that

specify rotations between pairs of dimensions. Each of the n-1 blocks contains n//2

pairs of independent rotations.

angles (torch.Tensor): A ``((n - 1) * n // 2,)`` shaped tensor containing all rotations

between pairs of dimensions.

blocks (torch.Tensor): A ``(n - 1, n // 2, 2)`` shaped tensor containing the indices that

specify rotations between pairs of dimensions. Each of the n-1 blocks contains n // 2

pairs of independent rotations.

Done, thanks.

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.

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

+        return U
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):


Missing return type hint.

I added the type hint as well as a longer explanation on what this return is.

thisac · 2025-12-30T00:52:24Z

+        U = self._create_rotation_matrix()
+        rotated_x = einsum(x, U, "... i, o i -> ... o")
+        if self.bias is not None:
+            rotated_x = rotated_x + self.bias


Suggested change

rotated_x = rotated_x + self.bias

rotated_x += self.bias

thisac · 2025-12-30T01:17:39Z

+from einops import einsum
+
+
+class NaiveGivensRotationLayer(nn.Module):


I'm not very keen on having a full on separate implementation here just to compare with/test the GivensRotationLayer. If this NaiveGivensRotationLayer is useful, should it be part of the package instead?

We discussed this in our one on one but, just for the record, there is no difference between the NaiveGivensRotationLayer and the GivensRotationLayer in the forward or backward passes. The naïve implementation is there to make sure that the forward and backward passes indeed match. The GivensRotationLayer should always be used because it has a substantially better runtime complexity. Thus, the naïve implementation is not useful—other than for a sanity check.

I think this class should go directly into the test file instead of creating a helper_models.py module. The naive module is only ever used in these tests.

I addressed this by movint the class to the test file

thisac · 2025-12-30T01:22:17Z

+    @parameterized.expand([(n, bias) for n in [4, 5, 6, 9, 10] for bias in [True, False]])
+    def test_forward_agreement(self, n, bias):


These tests do seem a bit too.. complex. Better to try and test more minimal aspects of the class, if possible. I'd much rather have separate integration-like tests that can assert that model behave as expected, while having these be strictly, small scale, isolated unit tests.

I added some tests to test invalid inputs too. These forward and backward tests are for testing that the correct input/output is given when compared to the naïve implementation. The model_probably_good test is done as unit test.

I added other unit tests where I test incorrect inputs as well. In ML models, the forward and backward passes should be what one expects them to be, and this module gives the opportunity to test this correctly. I do agree that we should separate other tests that (at least) I wrote, which have to do with training a model to see if the intended final trained state is what is expected. However, the tests I present in this PR are not the result of training but explicit comparisons with the naïve approach; I don't know if we could regard those as integration tests.

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.

anahitamansouri · 2026-01-02T22:37:02Z

+        self.n = n
+        self.n_angles = n * (n - 1) // 2
+        self.angles = nn.Parameter(torch.randn(self.n_angles))
+        blocks_edges = _get_blocks_edges(n)


You could directly return torch.LongTensor from get_blocks_edges to avoid the conversion.

I set _get_blocks_edges to a private function, so it shouldn't make a difference if I convert the list to a tensor in the orthogonal module or in the function itself.

I implemented your suggestion

kevinchern

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

kevinchern · 2026-01-05T19:26:17Z

+
+
+def _get_blocks_edges(n: int) -> list[list[tuple[int, int]]]:
+    """Uses the circle method for Round Robin pairing to create blocks of edges for parallel Givens


Suggested change

"""Uses the circle method for Round Robin pairing to create blocks of edges for parallel Givens

"""Uses the circle method for round-robin pairing to create blocks of edges for parallel Givens

(and other occurrences)

Should _get_blocks_edges should be a method in GivensRotation instead? The orthogonal module is general while this function is a helper function bespoke to GivensRotation.
cc @thisac

Maybe... though what would be the attribute of GivensRotation used in _get_blocks_edges? n only?

only n or as a @staticmethod

I took your suggestion

kevinchern · 2026-01-08T18:53:33Z

+        return grad_theta, None, None
+
+
+class GivensRotationLayer(nn.Module):


Can we rename to GivensRotation (parallel to nn.Linear)

Sounds good. Did that too

kevinchern · 2026-01-08T19:36:07Z

+    if n % 2 != 0:
+        n += 1  # Add a dummy dimension for odd n
+        is_odd = True
+    else:
+        is_odd = False


Could be cleaner like this 😛

Suggested change

if n % 2 != 0:

n += 1 # Add a dummy dimension for odd n

is_odd = True

else:

is_odd = False

odd = n % 2 != 0

if odd:

n += 1

or

odd = n % 2 ! = 0 n += odd

but this is less obvious.. (edit: not a big fan of n+=odd notation 😆)

It is cleaner! (the first suggestion, not the n+=odd 😆 )

kevinchern · 2026-01-08T19:39:10Z

+        ignored.
+    """
+    if n % 2 != 0:
+        n += 1  # Add a dummy dimension for odd n


Rule-of-thumb for comments: explain the "why" or motivation as opposed to "what" (which is clear in this context)

I will try to always adopt this rule of thumb!

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-01-08T20:52:13Z

+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor, None, None]:
+        """Computes the VJP needed for backward propagation.


Suggested change

"""Computes the VJP needed for backward propagation.

"""Computes the vector-Jacobian product needed for backward propagation.

kevinchern · 2026-01-08T23:03:15Z

+        idx_block = torch.arange(block_size, device=angles.device)
+        for b, block in enumerate(blocks):
+            # angles is of shape (n_angles,) containing all angles for contiguous blocks.
+            angles_in_block = angles[idx_block + b * blocks.size(1)]  # shape (n/2,)


Suggested change

angles_in_block = angles[idx_block + b * blocks.size(1)] # shape (n/2,)

angles_in_block = angles[idx_block + b * block_size] # shape (n/2,)

If I understand correctly, blocks.size(1) will be block_size

Ah yes, while writing the algorithm I though blocks could have different sizes if n is odd, but that is not true. All blocks will have the same block size.

kevinchern · 2026-01-08T23:23:10Z

+            c = torch.cos(angles_in_block)
+            s = torch.sin(angles_in_block)
+            i_idx = block[:, 0]
+            j_idx = block[:, 1]
+            r_i = c.unsqueeze(0) * U[:, i_idx] + s.unsqueeze(0) * U[:, j_idx]
+            r_j = -s.unsqueeze(0) * U[:, i_idx] + c.unsqueeze(0) * U[:, j_idx]


Unsqueeze once in the beginning

Suggested change

c = torch.cos(angles_in_block)

s = torch.sin(angles_in_block)

i_idx = block[:, 0]

j_idx = block[:, 1]

r_i = c.unsqueeze(0) * U[:, i_idx] + s.unsqueeze(0) * U[:, j_idx]

r_j = -s.unsqueeze(0) * U[:, i_idx] + c.unsqueeze(0) * U[:, j_idx]

c = torch.cos(angles_in_block).unsqueeze(0)

s = torch.sin(angles_in_block).unsqueeze(0)

i_idx = block[:, 0]

j_idx = block[:, 1]

r_i = c * U[:, i_idx] + s * U[:, j_idx]

r_j = -s * U[:, i_idx] + c * U[:, j_idx]

kevinchern · 2026-01-08T23:28:15Z

+        U = torch.eye(n, device=angles.device, dtype=angles.dtype)
+        block_size = n // 2
+        idx_block = torch.arange(block_size, device=angles.device)
+        for b, block in enumerate(blocks):


If we commit to using paper variable names here, we should be consistent and use, e.g., B instead of blocks.
If that's the case, I'd prefer to be a little more wasteful and have B = blocks to keep the input argument blocks instead of B. This inconsistency makes me lean towards named variables more (with a look-up table in the docstring).

Changed internally to use B instead.

kevinchern · 2026-01-08T23:29:49Z

+            r_i = c.unsqueeze(0) * U[:, i_idx] + s.unsqueeze(0) * U[:, j_idx]
+            r_j = -s.unsqueeze(0) * U[:, i_idx] + c.unsqueeze(0) * U[:, j_idx]


Are r_i and r_j are backwards?
I think it should be:

$\cos - \sin$ for i, and

$\sin + \cos$ for j.
Not sure if this has a significant impact on validity of method. If it does, then tests should be revised first to see why this error was not detected

Yes... well... in the paper the rotation matrices were written the other way around, I think. I did the math separately and this way everything is consistent.

Worth adding a # NOTE: here to highlight this distinction from paper, since we're also using variable names identical to paper

I added a .. note:: at the top of the class to explain that the paper performs rotations using the rows of U, but it is more standard to use the columns of U. It does not matter in the end because U is orthogonal, and using the rows or the columns is completely equivalent.

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

@VolodyaCO can you rebase on main?

kevinchern · 2026-04-21T16:17:14Z

+        Returns:
+            torch.Tensor: The nxn rotation matrix.
+        """
+        # Blocks is of shape (n_blocks, n/2, 2) containing indices for angles


Can you annotate with comments the corresponding equations from paper? This will make it easier to maintain and understand
e.g., ESS PR

some of the equations do not have equation numbers. I improved the backward pass code with in-line comments because it is the one that's actually difficult to follow with mappings from code variables to the variables used in the paper for improved readability.

do not have equation numbers

If it doesn't have number, then describe it unambiguously.

e.g., https://arxiv.org/pdf/2106.00003 (5) has 3 equations -> I'd just say # second equation of (5).

For another example
# the second unnumbered equation following (6)

A last example # the expression following the sentence "relevant storage of the gradient"

This was done mainly by copilot with AI. I went through the comments and fixed what was wrong (mainly that some references to the paper were incorrect in the sense that the location of the reference was wrongly stated, e.g. "after" instead of "before", or equation numbering wrong, e.g. "after equation (11)" instead of "after equation (12)".

kevinchern · 2026-04-21T16:24:06Z

+from einops import einsum
+
+
+class NaiveGivensRotationLayer(nn.Module):


I think this class should go directly into the test file instead of creating a helper_models.py module. The naive module is only ever used in these tests.

kevinchern · 2026-04-21T16:24:51Z

    def test_store_config(self):
        with self.subTest("Simple case"):
+
            class MyModel(torch.nn.Module):


VolodyaCO

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

VolodyaCO · 2026-05-04T22:39:37Z

    def test_store_config(self):
        with self.subTest("Simple case"):
+
            class MyModel(torch.nn.Module):


VolodyaCO · 2026-05-04T22:41:26Z

+        Returns:
+            torch.Tensor: The nxn rotation matrix.
+        """
+        # Blocks is of shape (n_blocks, n/2, 2) containing indices for angles


This was done mainly by copilot with AI. I went through the comments and fixed what was wrong (mainly that some references to the paper were incorrect in the sense that the location of the reference was wrongly stated, e.g. "after" instead of "before", or equation numbering wrong, e.g. "after equation (11)" instead of "after equation (12)".

thisac · 2026-05-05T22:50:52Z

+            tuple[torch.Tensor, None, None]: The gradient of the loss with respect to the input
+                angles. No calculation of gradients with respect to blocks or n is needed (cf.
+                forward method), so None is returned for these.


Returns shouldn't be indented. Same in other places in this file.

Suggested change

tuple[torch.Tensor, None, None]: The gradient of the loss with respect to the input

angles. No calculation of gradients with respect to blocks or n is needed (cf.

forward method), so None is returned for these.

tuple[torch.Tensor, None, None]: The gradient of the loss with respect to the input

angles. No calculation of gradients with respect to blocks or n is needed (cf.

forward method), so None is returned for these.

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

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

+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor, None, None]:
+        """Computes the vector-Jacobian product  needed for backward propagation.


Suggested change

"""Computes the vector-Jacobian product needed for backward propagation.

"""Computes the vector-Jacobian product needed for backward propagation.

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.

Address orthogonal module PR feedback.

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

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

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

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

thisac reviewed May 5, 2026

View reviewed changes

Add Givens orthogonal layer

9f31f15

Address orthogonal module PR feedback.

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

		from einops import einsum


		class NaiveGivensRotationLayer(nn.Module):

		@parameterized.expand([(n, bias) for n in [4, 5, 6, 9, 10] for bias in [True, False]])
		def test_forward_agreement(self, n, bias):



		def _get_blocks_edges(n: int) -> list[list[tuple[int, int]]]:
		"""Uses the circle method for Round Robin pairing to create blocks of edges for parallel Givens

	"""Uses the circle method for Round Robin pairing to create blocks of edges for parallel Givens
	"""Uses the circle method for round-robin pairing to create blocks of edges for parallel Givens

		return grad_theta, None, None


		class GivensRotationLayer(nn.Module):

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

	"""Computes the VJP needed for backward propagation.
	"""Computes the vector-Jacobian product needed for backward propagation.

	angles_in_block = angles[idx_block + b * blocks.size(1)] # shape (n/2,)
	angles_in_block = angles[idx_block + b * block_size] # shape (n/2,)

		r_i = c.unsqueeze(0) * U[:, i_idx] + s.unsqueeze(0) * U[:, j_idx]
		r_j = -s.unsqueeze(0) * U[:, i_idx] + c.unsqueeze(0) * U[:, j_idx]

		__all__ = ["GivensRotation"]


		class _RoundRobinGivens(torch.autograd.Function):

Conversation

VolodyaCO commented Dec 19, 2025

Uh oh!

VolodyaCO commented Dec 19, 2025

Uh oh!

Uh oh!

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

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 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!

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!

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!

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!

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!

Choose a reason for hiding this comment

Uh oh!

VolodyaCO commented Dec 31, 2025

Uh oh!

anahitamansouri left a comment

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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kevinchern Dec 19, 2025 •

edited

Loading

kevinchern commented Dec 19, 2025 •

edited

Loading