Add neural network modules, tests, and test utils

dwave/plugins/torch/nn.py

@@ -0,0 +1,127 @@
+# 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 __future__ import annotations
+
+import inspect
+from typing import TYPE_CHECKING, Callable
+
+if TYPE_CHECKING:
+    from functools import partial
+
+from functools import wraps
+from types import MappingProxyType
+
+import torch
+from torch import nn
+
+__all__ = ["store_config", "SkipLinear", "LinearBlock"]
+
+
+def store_config(fn: Callable) -> partial:
+    """A decorator that tracks and stores arguments of methods (excluding ``self``).
+
+    Args:
+        fn (Callable[object, ...]): A method whose arguments will be stored in ``self.config``.
+
+    Returns:
+        partial: Wrapper function that stores argument of method.
+    """
+    @wraps(fn)
+    def wrapper(self, *args, **kwargs):
+        """Store ``args``, ``kwargs``, and ``{"module_name": self.__class__.__name__}`` as a dictionary in ``self.config``.
+        """
+        sig = inspect.signature(fn)
+        bound = sig.bind(self, *args, **kwargs)
+        bound.apply_defaults()
+
+        config = {k: v for k, v in bound.arguments.items() if v != self}
+        config['module_name'] = self.__class__.__name__
+        self.config = MappingProxyType(config)
+
+        fn(self, *args, **kwargs)
+    return wrapper
+
+
+class SkipLinear(nn.Module):
+    """Applies a linear transformation to the incoming data: :math:`y = xA^T`.
+
+    This module is identity when ``din == dout``, otherwise, it is a linear transformation, i.e.,
+    no bias term.
+
+    Args:
+        din (int): Size of each input sample.
+        dout (int): Size of each output sample.
+    """
+
+    @store_config
+    def __init__(self, din: int, dout: int) -> None:
+        super().__init__()
+        if din == dout:
+            self.linear = nn.Identity()
+        else:
+            self.linear = nn.Linear(din, dout, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply a linear transformation to the input variable ``x``.
+
+        Args:
+            x (torch.Tensor): The input tensor.
+
+        Returns:
+            torch.Tensor: The linearly-transformed tensor of ``x``.
+        """
+        return self.linear(x)
+
+
+class LinearBlock(nn.Module):
+    """A linear block consisting of normalizations, linear transformations, dropout, relu, and a skip connection.
+
+    The module is composed of (in order):
+    1. a first layer norm,
+    2. a first linear transformation,
+    3. a dropout,
+    4. a relu activation,
+    5. a second layer norm,
+    6. a second linear layer, and, finally,
+    7. a skip connection from initial input to output.
+
+    Args:
+        din (int): Size of each input sample.
+        dout (int): Size of each output sample.
+        p (float): Dropout probability.
+    """
+    @store_config
+    def __init__(self, din: int, dout: int, p: float) -> None:
+        super().__init__()
+        self._skip = SkipLinear(din, dout)
+        self.block = nn.Sequential(
+            nn.LayerNorm(din),
+            nn.Linear(din, dout),
+            nn.Dropout(p),
+            nn.ReLU(),
+            nn.LayerNorm(dout),
+            nn.Linear(dout, dout),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Transforms the input `x` with the modules.
+
+        Args:
+            x (torch.Tensor): An input tensor.
+
+        Returns:
+            torch.Tensor: Output tensor.
+        """
+        return self.block(x) + self._skip(x)

releasenotes/notes/add-nn-modules-c29a092140eacbe1.yaml

@@ -0,0 +1,6 @@
+---
+features:
+  - Add the python module ``dwave.plugins.torch.nn`` containing commonly-used neural network modules
+    and patterns used to build more complex architectures.
+  - Add ``LinearBlock`` and ``SkipLinear` modules.
+  - Add utilities for testing torch modules added to the ``nn`` python submodule.

tests/test_nn.py

@@ -0,0 +1,52 @@
+import unittest
+
+import torch
+from parameterized import parameterized
+
+from dwave.plugins.torch.nn import LinearBlock, SkipLinear, store_config
+from tests.utils import model_probably_good
+
+
+class TestNN(unittest.TestCase):
+    """The tests in this class is, generally, concerned with two characteristics of the output.
+    1. Module outputs, probably, do not end with an activation function, and
+    2. the output tensor shapes are as expected.
+    """
+
+    def test_store_config(self):
+        # Check the Module stores configs as expected.
+        class MyModel(torch.nn.Module):
+            @store_config
+            def __init__(self, a, b=1, *, x=4, y='hello'):
+                super().__init__()
+
+            def forward(self, x):
+                return x
+        model = MyModel(a=123, x=5)
+        self.assertDictEqual(dict(model.config),
+                             {"a": 123, "b": 1, "x": 5, "y": "hello",
+                              "module_name": "MyModel"})
+
+    @parameterized.expand([0, 0.5, 1])
+    def test_LinearBlock(self, p):
+        din = 32
+        dout = 177
+        model = LinearBlock(din, dout, p)
+        self.assertTrue(model_probably_good(model, (din,), (dout,)))
+
+    def test_SkipLinear(self):
+        din = 33
+        dout = 99
+        model = SkipLinear(din, dout)
+        self.assertTrue(model_probably_good(model, (din,), (dout, )))
+        with self.subTest("Check identity for `din == dout`"):
+            dim = 123
+            model = SkipLinear(dim, dim)
+            x = torch.randn((dim,))
+            y = model(x)
+            self.assertTrue((x == y).all())
+            self.assertTrue(model_probably_good(model, (dim,), (dim, )))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_test_utils.py

@@ -0,0 +1,143 @@
+import unittest
+
+import torch
+from parameterized import parameterized
+
+from dwave.plugins.torch.nn import store_config
+from tests import utils
+
+
+class TestTestUtils(unittest.TestCase):
+
+    def test_probably_unconstrained(self):
+        x = torch.randn((1000, 10, 10))
+        self.assertTrue(utils._probably_unconstrained(x))
+
+        # Activate
+        self.assertFalse(utils._probably_unconstrained(x.sigmoid()))
+        self.assertFalse(utils._probably_unconstrained(x.relu()))
+        self.assertFalse(utils._probably_unconstrained(x.tanh()))
+
+    def test__are_all_spins(self):
+        # Scalar case
+        self.assertTrue(utils._are_all_spins(torch.tensor([1])))
+        self.assertTrue(utils._are_all_spins(torch.tensor([-1])))
+        self.assertFalse(utils._are_all_spins(torch.tensor([0])))
+
+        # Zeros
+        self.assertFalse(utils._are_all_spins(torch.tensor([0, 1])))
+        self.assertFalse(utils._are_all_spins(torch.tensor([0, -1])))
+        self.assertFalse(utils._are_all_spins(torch.tensor([0, 0])))
+        # Nonzeros
+        self.assertFalse(utils._are_all_spins(torch.tensor([1, 1.2])))
+        self.assertFalse(utils._are_all_spins(-torch.tensor([1, 1.2])))
+
+        # All spins
+        self.assertTrue(utils._are_all_spins(torch.tensor([-1, 1])))
+        self.assertTrue(utils._are_all_spins(torch.tensor([-1.0, 1.0])))
+
+    def test__has_zeros(self):
+        # Scalar
+        self.assertFalse(utils._has_zeros(torch.tensor([1])))
+        self.assertTrue(utils._has_zeros(torch.tensor([0])))
+        self.assertTrue(utils._has_zeros(torch.tensor([-0])))
+
+        # Tensor
+        self.assertTrue(utils._has_zeros(torch.tensor([0, 1])))
+
+    def test__has_mixed_signs(self):
+        # Single entries cannot have mixed signs
+        self.assertFalse(utils._has_mixed_signs(torch.tensor([-0])))
+        self.assertFalse(utils._has_mixed_signs(torch.tensor([0])))
+        self.assertFalse(utils._has_mixed_signs(torch.tensor([1])))
+        self.assertFalse(utils._has_mixed_signs(torch.tensor([-1])))
+
+        # Zeros are unsigned
+        self.assertFalse(utils._has_mixed_signs(torch.tensor([0, 0])))
+        self.assertFalse(utils._has_mixed_signs(torch.tensor([0, 1.2])))
+        self.assertFalse(utils._has_mixed_signs(torch.tensor([0, -1.2])))
+
+        # All entries have same sign
+        self.assertFalse(utils._has_mixed_signs(torch.tensor([0.4, 1.2])))
+        self.assertFalse(utils._has_mixed_signs(-torch.tensor([0.4, 1.2])))
+
+        # Finally!
+        self.assertTrue(utils._has_mixed_signs(torch.tensor([-0.1, 1.2])))
+
+    def test__bounded_in_plus_minus_one(self):
+        # Violation on one end
+        self.assertFalse(utils._bounded_in_plus_minus_one(torch.tensor([1.2])))
+        self.assertFalse(utils._bounded_in_plus_minus_one(torch.tensor([-1.2])))
+        self.assertFalse(utils._bounded_in_plus_minus_one(torch.tensor([1.2, 0])))
+        self.assertFalse(utils._bounded_in_plus_minus_one(torch.tensor([-1.2, 0])))
+
+        # Boundary
+        self.assertTrue(utils._bounded_in_plus_minus_one(torch.tensor([1])))
+        self.assertTrue(utils._bounded_in_plus_minus_one(torch.tensor([-1])))
+        self.assertTrue(utils._bounded_in_plus_minus_one(torch.tensor([1, -1])))
+        self.assertTrue(utils._bounded_in_plus_minus_one(torch.tensor([1, 0])))
+        self.assertTrue(utils._bounded_in_plus_minus_one(torch.tensor([0, 1])))
+
+        # Correct
+        self.assertTrue(utils._bounded_in_plus_minus_one(torch.tensor([0.5, 0.9, -0.2])))
+
+    @parameterized.expand([[dict(a=1, x=4)], [dict(a="hello")]])
+    def test__has_correct_config(self, kwargs):
+        class MyModel(torch.nn.Module):
+            @store_config
+            def __init__(self, a, b=2, *, x=4, y=5):
+                super().__init__()
+
+            def forward(self, x):
+                return torch.ones(5)
+        model = MyModel(**kwargs)
+        self.assertTrue(utils._has_correct_config(model))
+        self.assertFalse(utils._has_correct_config(torch.nn.Linear(5, 3)))
+
+    def test__shapes_match(self):
+        shape = (123, 456)
+        x = torch.randn(shape)
+        self.assertTrue(utils._shapes_match(x, shape))
+        self.assertFalse(utils._shapes_match(x, (1, 2, 3)))
+
+    def test_model_probably_good(self):
+        with self.subTest("Model should be good"):
+            class MyModel(torch.nn.Module):
+                @store_config
+                def __init__(self, a, b=2, *, x=4, y=5):
+                    super().__init__()
+
+                def forward(self, x):
+                    return 2*x
+            self.assertTrue(utils.model_probably_good(MyModel("hello"), (500, ), (500,)))
+
+        with self.subTest("Model should be bad: config not stored"):
+            class MyModel(torch.nn.Module):
+                def __init__(self, a, b=2, *, x=4, y=5):
+                    super().__init__()
+
+                def forward(self, x):
+                    return 2*x
+            self.assertFalse(utils.model_probably_good(MyModel("hello"), (500, ), (500,)))
+
+        with self.subTest("Model should be bad: shape mismatch"):
+            class MyModel(torch.nn.Module):
+                def __init__(self, a, b=2, *, x=4, y=5):
+                    super().__init__()
+
+                def forward(self, x):
+                    return torch.randn(500)
+            self.assertFalse(utils.model_probably_good(MyModel("hello"), (123, ), (123,)))
+
+        with self.subTest("Model should be bad: constrained output"):
+            class MyModel(torch.nn.Module):
+                def __init__(self, a, b=2, *, x=4, y=5):
+                    super().__init__()
+
+                def forward(self, x):
+                    return torch.ones_like(x)
+            self.assertFalse(utils.model_probably_good(MyModel("hello"), (123, ), (123,)))
+
+
+if __name__ == "__main__":
+    unittest.main()