torchnd 0.2.0__tar.gz → 0.3.0__tar.gz

{torchnd-0.2.0 → torchnd-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torchnd
-Version: 0.2.0
+Version: 0.3.0
 Summary: ND operations for PyTorch
 Author-email: Felix Zimmermann <fzimmermann89@gmail.com>
 License: GPL-3.0-or-later
@@ -85,13 +85,7 @@ out = conv_nd(x, weight, dim=(-2, -1), padding=1)
 ```python
 x = torch.randn(2, 4, 16, 16)
 weight = torch.randn(8, 4, 3, 3)
-out = conv_nd(
-    x, weight,
-    dim=(-2, -1),
-    stride=(2, 1),
-    padding=(1, 2),
-    dilation=(1, 2)
-)
+out = conv_nd(x, weight, dim=(-2, -1), stride=(2, 1), padding=(1, 2), dilation=(1, 2))
 ```
 
 **Modules:**
@@ -146,6 +140,19 @@ N-dimensional convolution with flexible dimension specification. Supports arbitr
 
 The `dim` parameter specifies which dimensions are spatial. The `channel_dim` parameter specifies the channel dimension. This allows for arbitrary tensor layouts (e.g., channel-last).
 
+`conv_nd`, `ConvNd`, and `ConvTransposeNd` accept `convolution=True` for
+mathematical convolution. The default keeps PyTorch's cross-correlation
+convention. Each call reverses the spatial kernel axes without changing the
+stored weight. Forward and transposed calls do not conjugate complex weights;
+`adjoint()` conjugates them and keeps the chosen convention. See the `conv_nd`
+docstring for the exact indexing and functional adjoint arguments.
+
+```python
+conv = ConvNd(4, 8, 3, padding=1, bias=False, convolution=True)
+y = conv(x)
+x_adjoint = conv.adjoint(y)
+```
+
 ### pad_nd
 
 N-dimensional padding and cropping with flexible dimension specification. Supports arbitrary dimension layouts, not limited to the last N dimensions. Negative padding values crop the tensor.

{torchnd-0.2.0 → torchnd-0.3.0}/README.md

@@ -54,13 +54,7 @@ out = conv_nd(x, weight, dim=(-2, -1), padding=1)
 ```python
 x = torch.randn(2, 4, 16, 16)
 weight = torch.randn(8, 4, 3, 3)
-out = conv_nd(
-    x, weight,
-    dim=(-2, -1),
-    stride=(2, 1),
-    padding=(1, 2),
-    dilation=(1, 2)
-)
+out = conv_nd(x, weight, dim=(-2, -1), stride=(2, 1), padding=(1, 2), dilation=(1, 2))
 ```
 
 **Modules:**
@@ -115,6 +109,19 @@ N-dimensional convolution with flexible dimension specification. Supports arbitr
 
 The `dim` parameter specifies which dimensions are spatial. The `channel_dim` parameter specifies the channel dimension. This allows for arbitrary tensor layouts (e.g., channel-last).
 
+`conv_nd`, `ConvNd`, and `ConvTransposeNd` accept `convolution=True` for
+mathematical convolution. The default keeps PyTorch's cross-correlation
+convention. Each call reverses the spatial kernel axes without changing the
+stored weight. Forward and transposed calls do not conjugate complex weights;
+`adjoint()` conjugates them and keeps the chosen convention. See the `conv_nd`
+docstring for the exact indexing and functional adjoint arguments.
+
+```python
+conv = ConvNd(4, 8, 3, padding=1, bias=False, convolution=True)
+y = conv(x)
+x_adjoint = conv.adjoint(y)
+```
+
 ### pad_nd
 
 N-dimensional padding and cropping with flexible dimension specification. Supports arbitrary dimension layouts, not limited to the last N dimensions. Negative padding values crop the tensor.

{torchnd-0.2.0 → torchnd-0.3.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "torchnd"
-version = "0.2.0"
+version = "0.3.0"
 description = "ND operations for PyTorch"
 readme = "README.md"
 requires-python = ">=3.10"

{torchnd-0.2.0 → torchnd-0.3.0}/src/torchnd/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.2.0"
+__version__ = "0.3.0"
 
 from torchnd.functional import (
     adjoint_linear_interpolation_nd,

{torchnd-0.2.0 → torchnd-0.3.0}/src/torchnd/functional.py

@@ -341,13 +341,16 @@ def conv_nd(
     groups: int = 1,
     transposed: bool = False,
     output_padding: int | tuple[int, ...] = 0,
+    convolution: bool = False,
 ) -> Tensor:
     """N-dimensional convolution with flexible dimension specification.
 
     Supports arbitrary spatial dimensions via recursive decomposition when
     native PyTorch implementations (1D, 2D, 3D) are unavailable.
 
-    Note: Strictly speaking, this implements correlation, not convolution. But it matches pytorch's naming convention.
+    By default, uses PyTorch's cross-correlation convention without conjugating
+    the kernel. Set ``convolution=True`` to reverse all spatial kernel axes
+    before applying the same operation. The stored/passed weight is unchanged.
 
     Parameters
     ----------
@@ -373,11 +376,25 @@ def conv_nd(
         If True, perform transposed (fractionally-strided) convolution.
     output_padding
         Additional size for transposed conv output (disambiguates shape).
+    convolution
+        If True, reverse the spatial kernel axes on each call. For one spatial
+        axis, the forward operation is
+        ``y[o, n] = sum(w[o, c, K-1-k] * x[c, n*s-p+k*d])``;
+        otherwise it uses ``w[o, c, k]``. Indices outside the input are zero.
+        Sums run over kernel indices and input channels within each group.
+        With ``transposed=True``, scatter ``x[c, n] * w[c, o, K-1-k]``
+        (or ``w[c, o, k]``) to ``y[o, n*s-p+k*d]``. Neither operation
+        conjugates the weight. To obtain the Hermitian adjoint, toggle
+        ``transposed`` and conjugate the weight, keeping ``convolution``
+        unchanged and choosing ``output_padding`` to recover the input shape.
 
     Returns
     -------
         Convolved output preserving input's dimension layout.
     """
+    if convolution:
+        weight = weight.flip(tuple(range(2, weight.ndim)))
+
     num_spatial = len(dim)
     stride = _normalize_tuple(stride, num_spatial)
     padding = _normalize_tuple(padding, num_spatial)

{torchnd-0.2.0 → torchnd-0.3.0}/src/torchnd/modules.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import math
-from abc import ABC, abstractmethod
+from abc import ABCMeta, abstractmethod
 
 import torch
 from torch import Tensor
@@ -12,7 +12,7 @@ from torch.utils.checkpoint import checkpoint
 from .functional import _normalize_tuple, adjoint_pad_nd, conv_nd, pad_nd
 
 
-class _ConvNdBase(torch.nn.Module, ABC):
+class _ConvNdBase(torch.nn.Module, metaclass=ABCMeta):
     """Base class for N-dimensional convolution modules."""
 
     transposed: bool = False
@@ -27,6 +27,7 @@ class _ConvNdBase(torch.nn.Module, ABC):
     dilation: tuple[int, ...]
     groups: int
     padding_mode: tuple[str, ...]
+    convolution: bool = False
 
     def __init__(
         self,
@@ -45,6 +46,7 @@ class _ConvNdBase(torch.nn.Module, ABC):
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
         weight: Tensor | None = None,
+        convolution: bool = False,
     ) -> None:
         """Initialize N-dimensional convolution.
 
@@ -83,6 +85,11 @@ class _ConvNdBase(torch.nn.Module, ABC):
         weight
             If provided, uses as external weight (enables hypernetwork usage).
             Registered as buffer to be included in state_dict.
+        convolution
+            If True, reverse all spatial kernel axes on each call to use
+            mathematical convolution. Defaults to PyTorch cross-correlation.
+            Forward never conjugates the weight; adjoint conjugates it and
+            keeps the same spatial reversal. See ``conv_nd`` for index formulas.
         """
         super().__init__()
 
@@ -102,6 +109,7 @@ class _ConvNdBase(torch.nn.Module, ABC):
         self.output_padding = _normalize_tuple(output_padding, num_spatial)
         self.dilation = _normalize_tuple(dilation, num_spatial)
         self.groups = groups
+        self.convolution = convolution
 
         if isinstance(padding_mode, str):
             self.padding_mode = (padding_mode,) * num_spatial
@@ -205,6 +213,8 @@ class _ConvNdBase(torch.nn.Module, ABC):
             parts.append(f"output_padding={self.output_padding}")
         if self.groups != 1:
             parts.append(f"groups={self.groups}")
+        if self.convolution:
+            parts.append("convolution=True")
         if self.bias is None:
             parts.append("bias=False")
         if not all(m == "zeros" for m in self.padding_mode):
@@ -261,6 +271,7 @@ class ConvNd(_ConvNdBase):
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
         weight: Tensor | None = None,
+        convolution: bool = False,
     ) -> None:
         """Initialize N-dimensional convolution.
 
@@ -295,6 +306,11 @@ class ConvNd(_ConvNdBase):
             Data type for parameters. Complex dtypes enable complex convolution.
         weight
             If provided, uses as external weight. Registered as buffer for state_dict.
+        convolution
+            If True, reverse all spatial kernel axes on each call to use
+            mathematical convolution. Defaults to PyTorch cross-correlation.
+            Forward never conjugates the weight; adjoint conjugates it and
+            keeps the same spatial reversal. See ``conv_nd`` for index formulas.
         """
         super().__init__(
             in_channels=in_channels,
@@ -312,6 +328,7 @@ class ConvNd(_ConvNdBase):
             device=device,
             dtype=dtype,
             weight=weight,
+            convolution=convolution,
         )
 
     def forward(self, x: Tensor) -> Tensor:
@@ -351,6 +368,7 @@ class ConvNd(_ConvNdBase):
             dilation=self.dilation,
             groups=self.groups,
             transposed=False,
+            convolution=self.convolution,
         )
 
         if self.bias is not None:
@@ -364,7 +382,8 @@ class ConvNd(_ConvNdBase):
         """Apply the adjoint convolution operator.
 
         The adjoint of a convolution is the transposed convolution with
-        conjugated weights (for complex) or identical weights (for real).
+        conjugated weights (for complex) or identical weights (for real),
+        retaining the same ``convolution`` setting.
 
         This method applies the adjoint using the same weights (and does not
         allocate a new module). Intended for iterative algorithms, not separate
@@ -436,6 +455,7 @@ class ConvNd(_ConvNdBase):
             dilation=self.dilation,
             groups=self.groups,
             transposed=True,
+            convolution=self.convolution,
             output_padding=out_pad,
         )
 
@@ -484,6 +504,7 @@ class ConvTransposeNd(_ConvNdBase):
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
         weight: Tensor | None = None,
+        convolution: bool = False,
     ) -> None:
         """Initialize N-dimensional transposed convolution.
 
@@ -520,6 +541,11 @@ class ConvTransposeNd(_ConvNdBase):
             Data type for parameters. Complex dtypes enable complex convolution.
         weight
             If provided, uses as external weight. Registered as buffer for state_dict.
+        convolution
+            If True, reverse all spatial kernel axes on each call to use
+            mathematical convolution. Defaults to PyTorch cross-correlation.
+            Forward never conjugates the weight; adjoint conjugates it and
+            keeps the same spatial reversal. See ``conv_nd`` for index formulas.
         """
         super().__init__(
             in_channels=in_channels,
@@ -537,6 +563,7 @@ class ConvTransposeNd(_ConvNdBase):
             device=device,
             dtype=dtype,
             weight=weight,
+            convolution=convolution,
         )
 
     def forward(self, x: Tensor) -> Tensor:
@@ -568,6 +595,7 @@ class ConvTransposeNd(_ConvNdBase):
             dilation=self.dilation,
             groups=self.groups,
             transposed=True,
+            convolution=self.convolution,
             output_padding=self.output_padding,
         )
 
@@ -602,7 +630,8 @@ class ConvTransposeNd(_ConvNdBase):
         """Apply the adjoint convolution operator.
 
         The adjoint of a transposed convolution is the (forward) convolution with
-        conjugated weights (for complex) or identical weights (for real).
+        conjugated weights (for complex) or identical weights (for real),
+        retaining the same ``convolution`` setting.
 
         This method applies the adjoint using the same weights (and does not
         allocate a new module). Intended for iterative algorithms, not separate
@@ -650,4 +679,5 @@ class ConvTransposeNd(_ConvNdBase):
             dilation=self.dilation,
             groups=self.groups,
             transposed=False,
+            convolution=self.convolution,
         )

{torchnd-0.2.0 → torchnd-0.3.0}/src/torchnd.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torchnd
-Version: 0.2.0
+Version: 0.3.0
 Summary: ND operations for PyTorch
 Author-email: Felix Zimmermann <fzimmermann89@gmail.com>
 License: GPL-3.0-or-later
@@ -85,13 +85,7 @@ out = conv_nd(x, weight, dim=(-2, -1), padding=1)
 ```python
 x = torch.randn(2, 4, 16, 16)
 weight = torch.randn(8, 4, 3, 3)
-out = conv_nd(
-    x, weight,
-    dim=(-2, -1),
-    stride=(2, 1),
-    padding=(1, 2),
-    dilation=(1, 2)
-)
+out = conv_nd(x, weight, dim=(-2, -1), stride=(2, 1), padding=(1, 2), dilation=(1, 2))
 ```
 
 **Modules:**
@@ -146,6 +140,19 @@ N-dimensional convolution with flexible dimension specification. Supports arbitr
 
 The `dim` parameter specifies which dimensions are spatial. The `channel_dim` parameter specifies the channel dimension. This allows for arbitrary tensor layouts (e.g., channel-last).
 
+`conv_nd`, `ConvNd`, and `ConvTransposeNd` accept `convolution=True` for
+mathematical convolution. The default keeps PyTorch's cross-correlation
+convention. Each call reverses the spatial kernel axes without changing the
+stored weight. Forward and transposed calls do not conjugate complex weights;
+`adjoint()` conjugates them and keeps the chosen convention. See the `conv_nd`
+docstring for the exact indexing and functional adjoint arguments.
+
+```python
+conv = ConvNd(4, 8, 3, padding=1, bias=False, convolution=True)
+y = conv(x)
+x_adjoint = conv.adjoint(y)
+```
+
 ### pad_nd
 
 N-dimensional padding and cropping with flexible dimension specification. Supports arbitrary dimension layouts, not limited to the last N dimensions. Negative padding values crop the tensor.

{torchnd-0.2.0 → torchnd-0.3.0}/src/torchnd.egg-info/SOURCES.txt

@@ -9,5 +9,7 @@ src/torchnd.egg-info/SOURCES.txt
 src/torchnd.egg-info/dependency_links.txt
 src/torchnd.egg-info/requires.txt
 src/torchnd.egg-info/top_level.txt
+tests/test_convolution_mode.py
 tests/test_functional.py
+tests/test_module_inheritance.py
 tests/test_modules.py

torchnd-0.3.0/tests/test_convolution_mode.py

@@ -0,0 +1,227 @@
+"""Independent values and adjoints for both kernel conventions."""
+
+from itertools import product
+
+import pytest
+import torch
+
+from torchnd import ConvNd, ConvTransposeNd, conv_nd
+
+
+def reference_conv(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    *,
+    convolution: bool,
+    transposed: bool,
+    stride: int = 1,
+    padding: int = 0,
+    dilation: int = 1,
+    groups: int = 1,
+    output_padding: int = 0,
+) -> torch.Tensor:
+    """Evaluate the defining sums directly, without any convolution backend."""
+    spatial = x.shape[2:]
+    kernel = weight.shape[2:]
+    if transposed:
+        out_shape = tuple(
+            (n - 1) * stride - 2 * padding + dilation * (k - 1) + 1 + output_padding for n, k in zip(spatial, kernel, strict=True)
+        )
+        out_channels = weight.shape[1] * groups
+    else:
+        out_shape = tuple((n + 2 * padding - dilation * (k - 1) - 1) // stride + 1 for n, k in zip(spatial, kernel, strict=True))
+        out_channels = weight.shape[0]
+    out = x.new_zeros(x.shape[0], out_channels, *out_shape)
+    in_per_group = x.shape[1] // groups
+    out_per_group = out_channels // groups
+    positions = spatial if transposed else out_shape
+    for batch, group, ci, co in product(range(x.shape[0]), range(groups), range(in_per_group), range(out_per_group)):
+        ic, oc = group * in_per_group + ci, group * out_per_group + co
+        for pos, offset in product(product(*(range(n) for n in positions)), product(*(range(k) for k in kernel))):
+            target = tuple(n * stride - padding + k * dilation for n, k in zip(pos, offset, strict=True))
+            limits = out_shape if transposed else spatial
+            if any(i < 0 or i >= n for i, n in zip(target, limits, strict=True)):
+                continue
+            wi = tuple(k - 1 - i for k, i in zip(kernel, offset, strict=True)) if convolution else offset
+            weight_index = (ic, co, *wi) if transposed else (oc, ci, *wi)
+            out_index = (batch, oc, *target) if transposed else (batch, oc, *pos)
+            in_index = (batch, ic, *pos) if transposed else (batch, ic, *target)
+            out[out_index] += x[in_index] * weight[weight_index]
+    return out
+
+
+@pytest.mark.parametrize("convolution", [False, True])
+@pytest.mark.parametrize("transposed", [False, True])
+@pytest.mark.parametrize("ndim", [0, 1, 2, 4])
+@pytest.mark.parametrize("dtype", [torch.float64, torch.complex128])
+def test_functional_reference(convolution: bool, transposed: bool, ndim: int, dtype: torch.dtype) -> None:
+    """Check grouped, strided, dilated kernels and channel-last layouts."""
+    generator = torch.Generator().manual_seed(3)
+    x = torch.randn(1, 4, *((3,) * ndim), dtype=dtype, generator=generator)
+    weight = torch.randn(4, 2, *((2,) * ndim), dtype=dtype, generator=generator)
+    expected = reference_conv(
+        x,
+        weight,
+        convolution=convolution,
+        transposed=transposed,
+        stride=2,
+        padding=1,
+        dilation=2,
+        groups=2,
+        output_padding=int(transposed),
+    )
+    actual = conv_nd(
+        x.moveaxis(1, -1),
+        weight,
+        dim=tuple(range(1, ndim + 1)),
+        channel_dim=-1,
+        stride=2,
+        padding=1,
+        dilation=2,
+        groups=2,
+        transposed=transposed,
+        output_padding=int(transposed),
+        convolution=convolution,
+    )
+    torch.testing.assert_close(actual.moveaxis(-1, 1), expected)
+    if not convolution:
+        default = conv_nd(
+            x.moveaxis(1, -1),
+            weight,
+            dim=tuple(range(1, ndim + 1)),
+            channel_dim=-1,
+            stride=2,
+            padding=1,
+            dilation=2,
+            groups=2,
+            transposed=transposed,
+            output_padding=int(transposed),
+        )
+        assert torch.equal(actual, default)
+
+
+@pytest.mark.parametrize("convolution", [False, True])
+@pytest.mark.parametrize("conv_type", [ConvNd, ConvTransposeNd])
+@pytest.mark.parametrize("ndim", [1, 4])
+@pytest.mark.parametrize("dtype", [torch.float64, torch.complex128])
+def test_module_reference_and_adjoint(
+    convolution: bool,
+    conv_type: type[ConvNd] | type[ConvTransposeNd],
+    ndim: int,
+    dtype: torch.dtype,
+) -> None:
+    """Modules and their adjoints agree with the independent matrix entries."""
+    generator = torch.Generator().manual_seed(5)
+    x = torch.randn(1, 4, *((3,) * ndim), dtype=dtype, generator=generator, requires_grad=True)
+    weight = torch.randn(4, 2, *((2,) * ndim), dtype=dtype, generator=generator)
+    module = conv_type(
+        4,
+        4,
+        2,
+        dim=tuple(range(2, ndim + 2)),
+        stride=2,
+        padding=1,
+        groups=2,
+        bias=False,
+        weight=weight,
+        convolution=convolution,
+    )
+    out = module(x)
+    expected = reference_conv(x, weight, convolution=convolution, transposed=module.transposed, stride=2, padding=1, groups=2)
+    torch.testing.assert_close(out, expected)
+    y = torch.randn(out.shape, dtype=dtype, generator=generator)
+    adjoint = module.adjoint(y, input_shape=(3,) * ndim) if isinstance(module, ConvNd) else module.adjoint(y)
+    (expected_adjoint,) = torch.autograd.grad(expected, x, y)
+    torch.testing.assert_close(adjoint, expected_adjoint)
+    torch.testing.assert_close(torch.vdot(out.flatten(), y.flatten()), torch.vdot(x.flatten(), adjoint.flatten()))
+
+
+@pytest.mark.parametrize("convolution", [False, True])
+@pytest.mark.parametrize("conv_type", [ConvNd, ConvTransposeNd])
+@pytest.mark.parametrize("padding_mode", ["zeros", "reflect", "replicate", "circular"])
+def test_padded_adjoint(
+    convolution: bool,
+    conv_type: type[ConvNd] | type[ConvTransposeNd],
+    padding_mode: str,
+) -> None:
+    """Boundary extension and its adjoint retain the chosen kernel convention."""
+    generator = torch.Generator().manual_seed(8)
+    module = conv_type(
+        2, 2, 3, dim=(-1,), padding=1, padding_mode=padding_mode, bias=False, dtype=torch.complex128, convolution=convolution
+    )
+    x = torch.randn(1, 2, 5, dtype=torch.complex128, generator=generator, requires_grad=True)
+    out = module(x)
+    y = torch.randn(out.shape, dtype=out.dtype, generator=generator)
+    (expected,) = torch.autograd.grad(out, x, y)
+    torch.testing.assert_close(module.adjoint(y), expected)
+
+
+@pytest.mark.parametrize("conv_type", [ConvNd, ConvTransposeNd])
+@pytest.mark.parametrize("external", ["internal", "parameter", "buffer"])
+@pytest.mark.parametrize("dtype", [torch.float64, torch.complex128])
+@pytest.mark.parametrize(
+    "device", ["cpu", pytest.param("cuda", marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA unavailable"))]
+)
+def test_weight_updates(conv_type: type[ConvNd] | type[ConvTransposeNd], external: str, dtype: torch.dtype, device: str) -> None:
+    """Each call uses the current weight and builds a fresh differentiable flip."""
+    weight = torch.arange(1.0, 4.0, device=device).to(dtype).unflatten(0, (1, 1, 3))
+    if dtype.is_complex:
+        weight = weight + 0.5j * weight.square()
+    if external == "parameter":
+        weight = torch.nn.Parameter(weight)
+    else:
+        weight.requires_grad_()
+    module = conv_type(
+        1,
+        1,
+        3,
+        dim=(-1,),
+        bias=False,
+        device=torch.device(device),
+        dtype=dtype,
+        weight=None if external == "internal" else weight,
+        convolution=True,
+    )
+    if external == "internal":
+        weight = module.weight
+    assert module.weight is weight
+    if external != "internal":
+        assert dict(module.named_buffers())["weight"] is weight
+    else:
+        assert dict(module.named_parameters())["weight"] is weight
+    x = torch.tensor([[[1.0, 2.0, 4.0, 8.0]]], dtype=dtype, device=device)
+    for step in range(2):
+        out = module(x)
+        reference_weight = weight.detach().clone().requires_grad_()
+        expected = reference_conv(x, reference_weight, convolution=True, transposed=module.transposed)
+        torch.testing.assert_close(out, expected)
+        out.abs().square().sum().backward()
+        (expected_grad,) = torch.autograd.grad(expected.abs().square().sum(), reference_weight)
+        torch.testing.assert_close(weight.grad, expected_grad)
+        y = torch.ones_like(out)
+        weight.grad = None
+        module.adjoint(y).abs().square().sum().backward()
+        expected_adjoint = reference_conv(y, reference_weight.conj(), convolution=True, transposed=not module.transposed)
+        (expected_grad,) = torch.autograd.grad(expected_adjoint.abs().square().sum(), reference_weight)
+        torch.testing.assert_close(weight.grad, expected_grad)
+        assert weight.grad is not None
+        assert torch.isfinite(weight.grad).all()
+        torch.testing.assert_close(module.state_dict()["weight"], weight)
+        with torch.no_grad():
+            weight.add_(step + 1)
+        weight.grad = None
+    assert "convolution=True" in repr(module)
+
+
+@pytest.mark.parametrize("conv_type", [ConvNd, ConvTransposeNd])
+def test_legacy_module_default(conv_type: type[ConvNd] | type[ConvTransposeNd]) -> None:
+    """Modules pickled before the setting existed keep cross-correlation."""
+    module = conv_type(1, 1, 3, dim=(-1,), bias=False, dtype=torch.float64)
+    x = torch.arange(5.0, dtype=torch.float64).unflatten(0, (1, 1, 5))
+    expected = module(x)
+    del module.convolution
+    assert module.convolution is False
+    torch.testing.assert_close(module(x), expected)
+    torch.testing.assert_close(
+        module.adjoint(expected), reference_conv(expected, module.weight, convolution=False, transposed=not module.transposed)
+    )

torchnd-0.3.0/tests/test_module_inheritance.py

@@ -0,0 +1,50 @@
+"""Convolution modules can share an abstract torch module base."""
+
+from abc import ABC, abstractmethod
+
+import pytest
+import torch
+
+from torchnd import ConvNd, ConvTransposeNd
+from torchnd.modules import _ConvNdBase
+
+
+class AbstractModule(ABC, torch.nn.Module):
+    """An abstract interface that places ABC before Module, as mr2 does."""
+
+    @abstractmethod
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply the module."""
+        raise NotImplementedError
+
+
+@pytest.mark.parametrize("conv_type", [ConvNd, ConvTransposeNd])
+@pytest.mark.parametrize("abstract_first", [False, True])
+def test_multiple_inheritance(conv_type: type[ConvNd] | type[ConvTransposeNd], abstract_first: bool) -> None:
+    """Either base order preserves construction, parameters, gradients, and the adjoint."""
+    bases = (AbstractModule, conv_type) if abstract_first else (conv_type, AbstractModule)
+    combined_type = type("CombinedConvolution", bases, {"forward": conv_type.forward})
+    module = combined_type(2, 3, 3, dim=(-1,), padding=1, bias=False, dtype=torch.float64)
+    assert isinstance(module, AbstractModule)
+    assert isinstance(module, conv_type)
+    assert dict(module.named_parameters())["weight"] is module.weight
+    torch.testing.assert_close(module.state_dict()["weight"], module.weight)
+
+    generator = torch.Generator().manual_seed(0)
+    x = torch.randn(1, 2, 5, generator=generator, dtype=torch.float64, requires_grad=True)
+    y = module(x)
+    assert y.shape == (1, 3, 5)
+    v = torch.randn(y.shape, generator=generator, dtype=y.dtype)
+    expected_adjoint = module.adjoint(v)
+    (gradient,) = torch.autograd.grad(y, x, grad_outputs=v, retain_graph=True)
+    torch.testing.assert_close(gradient, expected_adjoint)
+    y.square().sum().backward()
+    assert module.weight.grad is not None
+    assert module.float() is module
+    assert module.weight.dtype == torch.float32
+
+
+def test_base_remains_abstract() -> None:
+    """Using ABCMeta directly still prevents constructing the incomplete base class."""
+    with pytest.raises(TypeError, match="abstract"):
+        _ConvNdBase(1, 1, 3)  # type: ignore[abstract]

{torchnd-0.2.0 → torchnd-0.3.0}/tests/test_modules.py

@@ -122,7 +122,7 @@ class TestAdjoint:
     """Test adjoint operator correctness."""
 
     @pytest.mark.parametrize("ndim", [1, 2, 3])
-    @pytest.mark.parametrize("dtype", [torch.float32, torch.complex64])
+    @pytest.mark.parametrize("dtype", [torch.float64, torch.complex128])
     @pytest.mark.parametrize(
         "config",
         [
@@ -153,11 +153,10 @@ class TestAdjoint:
         lhs = torch.vdot(conv(x).flatten(), y.flatten())
         rhs = torch.vdot(x.flatten(), conv.adjoint(y, input_shape=input_shape[2:]).flatten())
 
-        rtol = 1e-4 if dtype == torch.float32 else 1e-5
-        torch.testing.assert_close(lhs, rhs, rtol=rtol, atol=rtol)
+        torch.testing.assert_close(lhs, rhs, rtol=1e-10, atol=1e-10)
 
     @pytest.mark.parametrize("ndim", [1, 2, 3])
-    @pytest.mark.parametrize("dtype", [torch.float32, torch.complex64])
+    @pytest.mark.parametrize("dtype", [torch.float64, torch.complex128])
     @pytest.mark.parametrize(
         "config",
         [
@@ -190,8 +189,7 @@ class TestAdjoint:
         lhs = torch.vdot(conv(x).flatten(), y.flatten())
         rhs = torch.vdot(x.flatten(), conv.adjoint(y).flatten())
 
-        rtol = 1e-4 if dtype == torch.float32 else 1e-5
-        torch.testing.assert_close(lhs, rhs, rtol=rtol, atol=rtol)
+        torch.testing.assert_close(lhs, rhs, rtol=1e-10, atol=1e-10)
 
     def test_raises_with_bias(self) -> None:
         """Adjoint raises ValueError if bias is present."""
@@ -460,16 +458,17 @@ class TestPaddingModeAdjoint:
             padding=2,
             padding_mode=mode,
             bias=False,
+            dtype=torch.float64,
         )
 
         input_shape = (x_shape[0], in_ch, *x_shape[2:])
-        x = torch.randn(input_shape)
+        x = torch.randn(input_shape, dtype=torch.float64)
         y = torch.randn_like(conv(x))
 
         lhs = torch.vdot(conv(x).flatten(), y.flatten())
         rhs = torch.vdot(x.flatten(), conv.adjoint(y, input_shape=input_shape[2:]).flatten())
 
-        torch.testing.assert_close(lhs, rhs, rtol=1e-4, atol=1e-5)
+        torch.testing.assert_close(lhs, rhs, rtol=1e-10, atol=1e-10)
 
     @pytest.mark.parametrize("mode", ["zeros", "reflect", "replicate", "circular"])
     @pytest.mark.parametrize("ndim", [1, 2])
@@ -486,18 +485,19 @@ class TestPaddingModeAdjoint:
             padding=2,
             padding_mode=mode,
             bias=False,
+            dtype=torch.float64,
         )
 
         spatial_in = tuple(s // 2 for s in x_shape[2:])
         input_shape = (x_shape[0], in_ch, *spatial_in)
 
-        x = torch.randn(input_shape)
+        x = torch.randn(input_shape, dtype=torch.float64)
         y = torch.randn_like(conv(x))
 
         lhs = torch.vdot(conv(x).flatten(), y.flatten())
         rhs = torch.vdot(x.flatten(), conv.adjoint(y).flatten())
 
-        torch.testing.assert_close(lhs, rhs, rtol=1e-4, atol=1e-5)
+        torch.testing.assert_close(lhs, rhs, rtol=1e-10, atol=1e-10)
 
     @pytest.mark.parametrize(
         "modes",
@@ -518,15 +518,16 @@ class TestPaddingModeAdjoint:
             padding=2,
             padding_mode=modes,
             bias=False,
+            dtype=torch.float64,
         )
 
-        x = torch.randn(2, 4, 16, 16)
+        x = torch.randn(2, 4, 16, 16, dtype=torch.float64)
         y = torch.randn_like(conv(x))
 
         lhs = torch.vdot(conv(x).flatten(), y.flatten())
         rhs = torch.vdot(x.flatten(), conv.adjoint(y, input_shape=(16, 16)).flatten())
 
-        torch.testing.assert_close(lhs, rhs, rtol=1e-4, atol=1e-5)
+        torch.testing.assert_close(lhs, rhs, rtol=1e-10, atol=1e-10)
 
     @pytest.mark.parametrize(
         "modes",
@@ -547,15 +548,16 @@ class TestPaddingModeAdjoint:
             padding=2,
             padding_mode=modes,
             bias=False,
+            dtype=torch.float64,
         )
 
-        x = torch.randn(2, 8, 8, 8)
+        x = torch.randn(2, 8, 8, 8, dtype=torch.float64)
         y = torch.randn_like(conv(x))
 
         lhs = torch.vdot(conv(x).flatten(), y.flatten())
         rhs = torch.vdot(x.flatten(), conv.adjoint(y).flatten())
 
-        torch.testing.assert_close(lhs, rhs, rtol=1e-4, atol=1e-5)
+        torch.testing.assert_close(lhs, rhs, rtol=1e-10, atol=1e-10)
 
     def test_3d_mixed_modes_adjoint(self) -> None:
         """Verify adjoint identity with 3 different padding modes."""
@@ -567,15 +569,16 @@ class TestPaddingModeAdjoint:
             padding=1,
             padding_mode=("reflect", "replicate", "circular"),
             bias=False,
+            dtype=torch.float64,
         )
 
-        x = torch.randn(2, 4, 8, 8, 8)
+        x = torch.randn(2, 4, 8, 8, 8, dtype=torch.float64)
         y = torch.randn_like(conv(x))
 
         lhs = torch.vdot(conv(x).flatten(), y.flatten())
         rhs = torch.vdot(x.flatten(), conv.adjoint(y, input_shape=(8, 8, 8)).flatten())
 
-        torch.testing.assert_close(lhs, rhs, rtol=1e-4, atol=1e-5)
+        torch.testing.assert_close(lhs, rhs, rtol=1e-10, atol=1e-10)
 
 
 class TestAdjointWithOutputPadding:
@@ -598,13 +601,14 @@ class TestAdjointWithOutputPadding:
             padding=1,
             output_padding=output_padding,
             bias=False,
+            dtype=torch.float64,
         )
 
         input_shape = (2, in_ch, *spatial_in)
-        x = torch.randn(input_shape)
+        x = torch.randn(input_shape, dtype=torch.float64)
         y = torch.randn_like(conv(x))
 
         lhs = torch.vdot(conv(x).flatten(), y.flatten())
         rhs = torch.vdot(x.flatten(), conv.adjoint(y).flatten())
 
-        torch.testing.assert_close(lhs, rhs, rtol=1e-4, atol=1e-5)
+        torch.testing.assert_close(lhs, rhs, rtol=1e-10, atol=1e-10)