Givens orthogonal layer

dwave/plugins/torch/nn/modules/__init__.py

@@ -14,4 +14,5 @@
 #
 
 from dwave.plugins.torch.nn.modules.linear import *
+from dwave.plugins.torch.nn.modules.orthogonal import *
 from dwave.plugins.torch.nn.modules.utils import *

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

@@ -0,0 +1,247 @@
+# Copyright 2025 D-Wave
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from collections import deque
+
+import torch
+import torch.nn as nn
+from einops import einsum
+
+from dwave.plugins.torch.nn.modules.utils import store_config
+
+__all__ = ["GivensRotation"]
+
+
+class _RoundRobinGivens(torch.autograd.Function):
+    """Implements custom forward and backward passes to implement the parallel algorithms in
+    https://arxiv.org/abs/2106.00003
+
+    .. note::
+        We adopt the notation from the paper, but instead of using the rows of U to compute
+        rotations, we follow the standard convention of using the columns of U. Since U is
+        orthogonal, this does not affect the result.
+    """
+
+    @staticmethod
+    def forward(ctx, angles: torch.Tensor, blocks: torch.Tensor, n: int) -> torch.Tensor:
+        """Creates a rotation matrix in n dimensions using parallel Givens transformations by
+        blocks.
+
+        Implements Algorithm 2 from https://arxiv.org/abs/2106.00003.
+
+        The algorithm reorders Givens rotations into n-1 blocks such that within each block,
+        all rotations operate on disjoint pairs of coordinates and thus commute. This enables
+        parallel computation within each block. See Section 3 ("Forward U Computation via
+        Round-Robin Sequences") for details on the round-robin sequence construction.
+
+        Args:
+            ctx (context): Stores information for backward propagation.
+            angles: A ``((n - 1) * n // 2,)`` shaped tensor containing all rotations between pairs
+                of dimensions.
+            blocks: 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.
+            n: Dimension of the space.
+
+        Returns:
+            The nxn rotation matrix.
+        """
+        # Blocks is of shape (n_blocks, n/2, 2) containing indices for angles
+        # Within each block, each Givens rotation is commuting, so we can apply them in parallel
+        U = torch.eye(n, device=angles.device, dtype=angles.dtype)
+        block_size = n // 2
+        idx_block = torch.arange(block_size, device=angles.device)
+        B = blocks  # to keep the same notation as in the paper
+        for b, block in enumerate(B):
+            # angles is of shape (n_angles,) containing all angles for contiguous blocks.
+            angles_in_block = angles[idx_block + b * block_size]  # shape (n/2,)
+            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]
+            U[:, i_idx] = r_i
+            U[:, j_idx] = r_j
+        ctx.save_for_backward(angles, B, U)
+        return U
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor, None, None]:
+        """Computes the vector-Jacobian product needed for backward propagation.
+
+        Implements Algorithm "Parallel JVP" from https://arxiv.org/abs/2106.00003
+        (presented in Section 4.2, "Computing the Gradient").
+
+        Args:
+            ctx (context): Contains information for backward propagation.
+            grad_output: A tensor containing the partial derivatives for the loss with respect to
+                the output of the forward pass, i.e., dL/dU.
+
+        Returns:
+            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.
+        """
+        angles, B, Ufwd_saved = ctx.saved_tensors
+        # U^fwd represents U^{1:k-1} at block k (Eq. (11)). Algorithm 3 initializes U^fwd <- U.
+        Ufwd = Ufwd_saved.clone()
+
+        # Gamma = dL/dU. M is defined in Eq. (15) as M = U^bck @ Gamma^T. Algorithm 3 initializes
+        # M <- Gamma^T.
+        M = grad_output.T
+
+        n = M.size(1)
+        block_size = n // 2
+        # A is the temporary n/2 x n matrix from Section 4.2 (text before Algorithm 3).
+        # Each row m corresponds to one edge e = (i, j) in the current block via mapping m(e).
+        A = torch.zeros((block_size, n), device=angles.device, dtype=angles.dtype)
+
+        grad_theta = torch.zeros_like(angles, dtype=angles.dtype)  # d = dL/dθ (to be computed)
+        idx_block = torch.arange(block_size, device=angles.device)
+        for b, block in enumerate(B):
+            i_idx = block[:, 0]
+            j_idx = block[:, 1]
+            # θ_e for the current round-robin block b_k (contiguous angles for this block)
+            angles_in_block = angles[idx_block + b * block_size]  # shape (n/2,)
+
+            c = torch.cos(angles_in_block)
+            s = torch.sin(angles_in_block)
+
+            # Algorithm 3 first loop (see alsoSection 4.1, around Eq. (12)).
+            r_i = c.unsqueeze(1) * Ufwd[i_idx] + s.unsqueeze(1) * Ufwd[j_idx]
+            r_j = -s.unsqueeze(1) * Ufwd[i_idx] + c.unsqueeze(1) * Ufwd[j_idx]
+            Ufwd[i_idx] = r_i
+            Ufwd[j_idx] = r_j
+
+            # Algorithm 3 second loop, see also Eq. (15).
+            r_i = c.unsqueeze(0) * M[:, i_idx] + s.unsqueeze(0) * M[:, j_idx]
+            r_j = -s.unsqueeze(0) * M[:, i_idx] + c.unsqueeze(0) * M[:, j_idx]
+            M[:, i_idx] = r_i
+            M[:, j_idx] = r_j
+
+            # Algorithm 3 third loop (see also Section 4.2, text after Eq. (16)).
+            A[:] = M[:, j_idx].T * Ufwd[i_idx] - M[:, i_idx].T * Ufwd[j_idx]
+
+            # Algorithm 3 reduction (see also Section 4.2,  before Algorithm 3)
+            grad_theta[idx_block + b * block_size] = A.sum(dim=1)
+        return grad_theta, None, None
+
+
+class GivensRotation(nn.Module):
+    """An orthogonal layer implementing a rotation using a sequence of Givens rotations arranged in
+    a round-robin fashion.
+
+    Angles are arranged into blocks, where each block references rotations that can be applied in
+    parallel because these rotations commute.
+
+    Args:
+        n: Dimension of the input and output space. Must be at least 2.
+        bias: If True, adds a learnable bias to the output. Default: True.
+    """
+
+    @store_config
+    def __init__(self, n: int, bias: bool = True):
+        super().__init__()
+        if not isinstance(n, int) or n <= 1:
+            raise ValueError(f"n must be an integer greater than 1, {n} was passed")
+        if not isinstance(bias, bool):
+            raise ValueError(f"bias must be a boolean, {bias} was passed")
+        self._n = n
+        self._n_angles = n * (n - 1) // 2
+        self.angles = nn.Parameter(torch.randn(self._n_angles))
+        blocks_edges = self._get_blocks_edges(n)
+        self.register_buffer("_blocks", blocks_edges)
+        if bias:
+            self.bias = nn.Parameter(torch.zeros(n))
+        else:
+            self.register_parameter("bias", None)
+
+    @property
+    def n(self) -> int:
+        """Returns the dimension of the input and output space."""
+        return self._n
+
+    @property
+    def n_angles(self) -> int:
+        """Returns the number of angles used in the Givens rotation."""
+        return self._n_angles
+
+    @property
+    def blocks(self) -> torch.Tensor:
+        """Returns the blocks of edges used in the Givens rotation."""
+        return self._blocks
+
+    @staticmethod
+    def _get_blocks_edges(n: int) -> torch.Tensor:
+        """Uses the circle method for Round Robin pairing to create blocks of edges for parallel
+        Givens rotations.
+
+        A block is a list of pairs of indices indicating which coordinates to rotate together. Pairs
+        in the same block can be rotated in parallel since they commute.
+
+        Args:
+            n: Dimension of the vector space onto which an orthogonal layer will be built.
+
+        Returns:
+            Blocks of edges for parallel Givens rotations stored in a tensor of shape
+            ``(n - 1, n // 2, 2)``.
+
+        .. note::
+            If n is odd, a dummy dimension is added to make it even. When using the resulting blocks
+            to build an orthogonal transformation, rotations involving the dummy dimension should be
+            ignored.
+        """
+        is_odd = bool(n % 2 != 0)
+        if is_odd:
+            # The circle method requires an even number of nodes, so we add a dummy dimension, the
+            # additional rotations involving this dimension will be ignored later.
+            n += 1
+
+        def circle_method(sequence):
+            seq_first_half = sequence[: len(sequence) // 2]
+            seq_second_half = sequence[len(sequence) // 2 :][::-1]
+            return list(zip(seq_first_half, seq_second_half))
+
+        blocks = []
+        sequence = list(range(n))
+        sequence_deque = deque(sequence[1:])
+        for _ in range(n - 1):
+            pairs = circle_method(sequence)
+            if is_odd:
+                # Remove pairs involving the dummy dimension:
+                pairs = [pair for pair in pairs if n - 1 not in pair]
+            blocks.append(pairs)
+            sequence_deque.rotate(1)
+            sequence[1:] = list(sequence_deque)
+        return torch.tensor(blocks, dtype=torch.long)
+
+    def _create_rotation_matrix(self) -> torch.Tensor:
+        """Computes the Givens rotation matrix."""
+        return _RoundRobinGivens.apply(self.angles, self._blocks, self._n)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Applies the Givens rotation to the input tensor ``x``.
+
+        Args:
+            x: Input tensor of shape ``(..., n)``.
+
+        Returns:
+            Rotated tensor of shape ``(..., n)``.
+        """
+        unitary = self._create_rotation_matrix()
+        rotated_x = einsum(x, unitary, "... i, o i -> ... o")
+        if self.bias is not None:
+            rotated_x += self.bias
+        return rotated_x

pyproject.toml

@@ -32,6 +32,7 @@ dependencies = [
     "dimod",
     "dwave-system",
     "dwave-hybrid",
+    "einops",
 ]
 
 [project.readme]

releasenotes/notes/add-orthogonal-rotation-givens-layer-320f2cb6046ac1da.yaml

@@ -0,0 +1,4 @@
+---
+features:
+  - |
+    Add orthogonal rotation layer using Givens rotations.

requirements.txt

@@ -6,6 +6,7 @@ torch==2.9.1
 dimod==0.12.21
 dwave-system==1.34.0
 dwave-hybrid==0.6.14
+einops==0.8.1
 
 # Development requirements
 reno==4.1.0