torchnd 0.1.0__py3-none-any.whl

torchnd/__init__.py

@@ -0,0 +1,14 @@
+__version__ = "0.1.0"
+
+from torchnd.functional import adjoint_pad_nd, conv_nd, pad_nd, pad_or_crop_to_size
+from torchnd.modules import ConvNd, ConvTransposeNd
+
+__all__ = [
+    "adjoint_pad_nd",
+    "conv_nd",
+    "pad_nd",
+    "pad_or_crop_to_size",
+    "ConvNd",
+    "ConvTransposeNd",
+    "__version__",
+]

torchnd/functional.py

@@ -0,0 +1,478 @@
+"""N-dimensional convolution with flexible dimension specification."""
+
+from collections.abc import Callable, Sequence
+
+import torch
+from einops import rearrange
+from torch import Tensor
+
+CONV_REGISTRY: dict[int, Callable[..., Tensor]] = {
+    1: torch.nn.functional.conv1d,
+    2: torch.nn.functional.conv2d,
+    3: torch.nn.functional.conv3d,
+}
+
+CONV_TRANSPOSE_REGISTRY: dict[int, Callable[..., Tensor]] = {
+    1: torch.nn.functional.conv_transpose1d,
+    2: torch.nn.functional.conv_transpose2d,
+    3: torch.nn.functional.conv_transpose3d,
+}
+
+
+def _normalize_tuple(value: int | tuple[int, ...], length: int) -> tuple[int, ...]:
+    return (value,) * length if isinstance(value, int) else tuple(value)
+
+
+def conv_nd(
+    x: Tensor,
+    weight: Tensor,
+    dim: tuple[int, ...] = (-2, -1),
+    channel_dim: int = 1,
+    stride: int | tuple[int, ...] = 1,
+    padding: int | tuple[int, ...] = 0,
+    dilation: int | tuple[int, ...] = 1,
+    groups: int = 1,
+    transposed: bool = False,
+    output_padding: int | tuple[int, ...] = 0,
+) -> 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.
+
+    Parameters
+    ----------
+    x
+        Input tensor with batch, channel, and spatial dimensions.
+    weight
+        Convolution kernel.
+        Forward: ``(C_out, C_in // groups, *kernel_sizes)``.
+        Transposed: ``(C_in, C_out // groups, *kernel_sizes)``.
+    dim
+        Spatial dimensions to convolve over.
+    channel_dim
+        Channel dimension index.
+    stride
+        Stride of the convolution; controls the step size the kernel moves along each spatial dimension. Can be a single integer or a tuple giving a value per dimension.
+    padding
+        Padding added to each spatial dimension. Can be a single integer or a tuple specifying padding for each dimension.
+    dilation
+        Spacing between kernel elements along each spatial dimension. Can be a single integer or per-dimension tuple.
+    groups
+        Grouped convolution groups.
+    transposed
+        If True, perform transposed (fractionally-strided) convolution.
+    output_padding
+        Additional size for transposed conv output (disambiguates shape).
+
+    Returns
+    -------
+        Convolved output preserving input's dimension layout.
+    """
+    num_spatial = len(dim)
+    stride = _normalize_tuple(stride, num_spatial)
+    padding = _normalize_tuple(padding, num_spatial)
+    dilation = _normalize_tuple(dilation, num_spatial)
+    output_padding = _normalize_tuple(output_padding, num_spatial)
+
+    def conv(a: Tensor, b: Tensor) -> Tensor:
+        return _dispatch(a, b, dim, channel_dim, stride, padding, dilation, groups, transposed, output_padding)
+
+    match (x.is_complex(), weight.is_complex()):
+        case (True, True):
+            return torch.complex(
+                conv(x.real, weight.real) - conv(x.imag, weight.imag),
+                conv(x.real, weight.imag) + conv(x.imag, weight.real),
+            )
+        case (True, False):
+            return torch.complex(conv(x.real, weight), conv(x.imag, weight))
+        case (False, True):
+            return torch.complex(conv(x, weight.real), conv(x, weight.imag))
+        case _:
+            return conv(x, weight)
+
+
+def _dispatch(
+    x: Tensor,
+    weight: Tensor,
+    dim: tuple[int, ...],
+    channel_dim: int,
+    stride: tuple[int, ...],
+    padding: tuple[int, ...],
+    dilation: tuple[int, ...],
+    groups: int,
+    transposed: bool,
+    output_padding: tuple[int, ...],
+) -> Tensor:
+    """Permute to standard layout, convolve, permute back."""
+    ndim = x.ndim
+    spatial_dims = tuple(d % ndim for d in dim)
+    channel_dim_normalized = channel_dim % ndim
+    batch_dims = tuple(d for d in range(ndim) if d not in set(spatial_dims) and d != channel_dim_normalized)
+
+    to_standard = batch_dims + (channel_dim_normalized,) + spatial_dims
+    x_perm = x.permute(to_standard)
+
+    num_batch = len(batch_dims)
+    batch_shape = x_perm.shape[:num_batch]
+    x_flat = x_perm.reshape(-1 if batch_shape else 1, *x_perm.shape[num_batch:])
+
+    out_flat = _conv_core(x_flat, weight, stride, padding, dilation, groups, transposed, output_padding)
+
+    out = out_flat.reshape(*batch_shape, *out_flat.shape[1:]) if batch_shape else out_flat.squeeze(0)
+    from_standard = torch.argsort(torch.tensor(to_standard)).tolist()
+    return out.permute(from_standard)
+
+
+def _conv_core(
+    x: Tensor,
+    weight: Tensor,
+    stride: tuple[int, ...],
+    padding: tuple[int, ...],
+    dilation: tuple[int, ...],
+    groups: int,
+    transposed: bool,
+    output_padding: tuple[int, ...],
+) -> Tensor:
+    """Dispatch to native implementation or recursive fallback."""
+    num_spatial = len(stride)
+    registry = CONV_TRANSPOSE_REGISTRY if transposed else CONV_REGISTRY
+
+    if num_spatial == 0:
+        return registry[1](x.unsqueeze(-1), weight.unsqueeze(-1), groups=groups).squeeze(-1)
+
+    if num_spatial in registry:
+        kwargs = {
+            "bias": None,
+            "stride": stride,
+            "padding": padding,
+            "dilation": dilation,
+            "groups": groups,
+        }
+        if transposed:
+            kwargs["output_padding"] = output_padding
+        return registry[num_spatial](x, weight, **kwargs)
+
+    if transposed:
+        return _conv_transpose_recursive(x, weight, stride, padding, dilation, groups, output_padding)
+    return _conv_recursive(x, weight, stride, padding, dilation, groups)
+
+
+def _conv_recursive(
+    x: Tensor,
+    weight: Tensor,
+    stride: tuple[int, ...],
+    padding: tuple[int, ...],
+    dilation: tuple[int, ...],
+    groups: int,
+) -> Tensor:
+    """Decompose N-D conv: sum over first kernel dim of (N-1)-D convs."""
+    if padding[0] > 0:
+        pad_spec = (0, 0) * (x.ndim - 3) + (padding[0], padding[0])
+        x = torch.nn.functional.pad(x, pad_spec)
+
+    kernel_size = weight.shape[2]
+    input_size = x.shape[2]
+    batch_size = x.shape[0]
+    output_size = (input_size - dilation[0] * (kernel_size - 1) - 1) // stride[0] + 1
+
+    accumulator: Tensor | None = None
+    for ki in range(kernel_size):
+        start = ki * dilation[0]
+        x_slice = x[:, :, start : start + (output_size - 1) * stride[0] + 1 : stride[0]]
+
+        out = _conv_core(
+            rearrange(x_slice, "n c d ... -> (n d) c ..."),
+            weight[:, :, ki],
+            stride[1:],
+            padding[1:],
+            dilation[1:],
+            groups,
+            False,
+            (),
+        )
+        out = rearrange(out, "(n d) c ... -> n c d ...", n=batch_size, d=output_size)
+        accumulator = out if accumulator is None else accumulator + out
+
+    assert accumulator is not None
+    return accumulator
+
+
+def _conv_transpose_recursive(
+    x: Tensor,
+    weight: Tensor,
+    stride: tuple[int, ...],
+    padding: tuple[int, ...],
+    dilation: tuple[int, ...],
+    groups: int,
+    output_padding: tuple[int, ...],
+) -> Tensor:
+    """Decompose N-D transpose conv: scatter (N-1)-D results into output."""
+    kernel_size = weight.shape[2]
+    input_size = x.shape[2]
+    batch_size = x.shape[0]
+    out_channels = weight.shape[1] * groups
+    full_size = (input_size - 1) * stride[0] + dilation[0] * (kernel_size - 1) + 1
+
+    result: Tensor | None = None
+    for ki in range(kernel_size):
+        out = _conv_core(
+            rearrange(x, "n c d ... -> (n d) c ..."),
+            weight[:, :, ki],
+            stride[1:],
+            padding[1:],
+            dilation[1:],
+            groups,
+            True,
+            output_padding[1:],
+        )
+        out = rearrange(out, "(n d) c ... -> n c d ...", n=batch_size, d=input_size)
+
+        if result is None:
+            result = x.new_zeros(batch_size, out_channels, full_size, *out.shape[3:])
+
+        start = ki * dilation[0]
+        end = start + input_size * stride[0]
+        result[:, :, start : end : stride[0]] = result[:, :, start : end : stride[0]] + out
+
+    assert result is not None
+
+    if padding[0] > 0:
+        result = result[:, :, padding[0] :]
+
+    right_adjust = padding[0] - output_padding[0]
+    if right_adjust > 0:
+        result = result[:, :, :-right_adjust]
+    elif right_adjust < 0:
+        pad_spec = (0, 0) * (result.ndim - 3) + (0, -right_adjust)
+        result = torch.nn.functional.pad(result, pad_spec)
+
+    return result
+
+
+def pad_nd(
+    input: Tensor,
+    pad: Sequence[int],
+    mode: str = "constant",
+    value: float = 0.0,
+    dims: Sequence[int] | None = None,
+) -> Tensor:
+    """
+    N-dimensional padding/cropping supporting arbitrary dimensions.
+
+    Parameters
+    ----------
+    input
+        Input tensor.
+    pad
+        Padding sizes as pairs (left, right) per dimension.
+        Negative values crop. Without `dims`: PyTorch convention.
+    mode
+        'constant', 'circular', 'reflect', or 'replicate'.
+    value
+        Fill value for 'constant' mode.
+    dims
+        Dimensions to pad. If None, pads last N dimensions.
+
+    Returns
+    -------
+        Padded/cropped tensor.
+    """
+    if len(pad) % 2 != 0:
+        raise ValueError("pad must have even length")
+
+    n_dims = len(pad) // 2
+    if n_dims == 0:
+        return input
+
+    ndim = input.ndim
+
+    if dims is None:
+        target_dims = list(range(ndim - 1, ndim - n_dims - 1, -1))
+        pad_pairs = [(pad[2 * i], pad[2 * i + 1]) for i in range(n_dims)]
+    else:
+        if len(dims) != n_dims:
+            raise ValueError("len(dims) must equal len(pad) // 2")
+        target_dims = [d % ndim for d in dims]
+        pad_pairs = [(pad[2 * i], pad[2 * i + 1]) for i in range(n_dims)]
+        if len(set(target_dims)) != len(target_dims):
+            raise ValueError("duplicate dimensions not supported")
+
+    if mode == "constant":
+        return _pad_constant(input, target_dims, pad_pairs, value)
+
+    return _pad_non_constant(input, target_dims, pad_pairs, mode)
+
+
+def pad_or_crop_to_size(
+    input: Tensor,
+    size: Sequence[int],
+    mode: str = "constant",
+    value: float = 0.0,
+    dims: Sequence[int] | None = None,
+) -> Tensor:
+    """
+    Pad or crop tensor to target size, centering the original content.
+
+    Parameters
+    ----------
+    input
+        Input tensor.
+    size
+        Target sizes for each dimension to adjust.
+    mode
+        'constant', 'circular', 'reflect', or 'replicate'.
+    value
+        Fill value for 'constant' mode.
+    dims
+        Dimensions to adjust. If None, adjusts last len(size) dimensions.
+
+    Returns
+    -------
+        Tensor with specified dimensions resized to target.
+    """
+    ndim = input.ndim
+
+    if dims is None:
+        target_dims = list(range(ndim - len(size), ndim))
+    else:
+        if len(dims) != len(size):
+            raise ValueError("len(dims) must equal len(size)")
+        target_dims = [d % ndim for d in dims]
+
+    pad_list: list[int] = []
+    for dim, target in zip(target_dims, size):
+        diff = target - input.shape[dim]
+        pad_list.extend([diff // 2, diff - diff // 2])
+
+    return pad_nd(input, pad_list, mode=mode, value=value, dims=target_dims)
+
+
+def _pad_constant(
+    input: Tensor,
+    target_dims: list[int],
+    pad_pairs: list[tuple[int, int]],
+    value: float,
+) -> Tensor:
+    dim_to_pad = dict(zip(target_dims, pad_pairs))
+
+    full_pad: list[int] = []
+    for d in range(input.ndim - 1, -1, -1):
+        full_pad.extend(dim_to_pad.get(d, (0, 0)))
+
+    while len(full_pad) > 2 and full_pad[-2:] == [0, 0]:
+        full_pad = full_pad[:-2]
+
+    return torch.nn.functional.pad(input, full_pad, mode="constant", value=value)
+
+
+def _pad_non_constant(
+    input: Tensor,
+    target_dims: list[int],
+    pad_pairs: list[tuple[int, int]],
+    mode: str,
+) -> Tensor:
+    result = input
+
+    for i in range(0, len(target_dims), 3):
+        chunk_dims = target_dims[i : i + 3]
+        chunk_pads = pad_pairs[i : i + 3]
+
+        other_dims = [d for d in range(result.ndim) if d not in chunk_dims]
+        perm = other_dims + chunk_dims
+
+        inv_perm = [0] * len(perm)
+        for old_pos, new_pos in enumerate(perm):
+            inv_perm[new_pos] = old_pos
+
+        result = result.permute(perm)
+        batch_shape = result.shape[: len(other_dims)]
+
+        if len(batch_shape) == 0:
+            result = result.unsqueeze(0)
+        elif len(batch_shape) > 1:
+            result = result.flatten(0, len(batch_shape) - 1)
+
+        n_unsqueeze = max(0, len(chunk_dims) + 2 - result.ndim)
+        for _ in range(n_unsqueeze):
+            result = result.unsqueeze(1)
+
+        flat_pad = [p for lr in reversed(chunk_pads) for p in lr]
+        result = torch.nn.functional.pad(result, flat_pad, mode=mode)
+
+        for _ in range(n_unsqueeze):
+            result = result.squeeze(1)
+
+        if len(batch_shape) == 0:
+            result = result.squeeze(0)
+        elif len(batch_shape) > 1:
+            result = result.unflatten(0, batch_shape)
+
+        result = result.permute(inv_perm)
+
+    return result
+
+
+def adjoint_pad_nd(
+    input: Tensor,
+    pad: Sequence[int],
+    mode: str = "constant",
+    value: float = 0.0,
+    dims: Sequence[int] | None = None,
+    original_sizes: Sequence[int] | None = None,
+) -> Tensor:
+    """
+    Adjoint of pad_nd via autograd - computes backward of padding as forward op.
+
+    Parameters
+    ----------
+    input
+        Padded tensor (output space of pad_nd).
+    pad
+        Padding sizes as pairs (left, right) per dimension (same as pad_nd).
+    mode
+        'constant', 'circular', 'reflect', or 'replicate'.
+    value
+        Fill value for 'constant' mode.
+    dims
+        Dimensions that were padded. If None, assumes last N dimensions.
+    original_sizes
+        Original sizes of the padded dimensions before padding.
+        Required to reconstruct the unpadded shape.
+
+    Returns
+    -------
+        Tensor with padding adjoint applied (unpadded shape).
+    """
+    if len(pad) % 2 != 0:
+        raise ValueError("pad must have even length")
+
+    n_dims = len(pad) // 2
+    if n_dims == 0:
+        return input
+
+    ndim = input.ndim
+
+    if dims is None:
+        target_dims = list(range(ndim - 1, ndim - n_dims - 1, -1))
+    else:
+        if len(dims) != n_dims:
+            raise ValueError("len(dims) must equal len(pad) // 2")
+        target_dims = [d % ndim for d in dims]
+
+    if original_sizes is None:
+        original_sizes = []
+        for i, d in enumerate(target_dims):
+            left, right = pad[2 * i], pad[2 * i + 1]
+            original_sizes.append(input.shape[d] - left - right)
+
+    unpadded_shape = list(input.shape)
+    for i, d in enumerate(target_dims):
+        unpadded_shape[d] = original_sizes[i]
+
+    x = torch.zeros(unpadded_shape, device=input.device, dtype=input.dtype, requires_grad=True)
+    y = pad_nd(x, pad, mode=mode, value=value, dims=dims)
+    (gx,) = torch.autograd.grad(y, x, grad_outputs=input, retain_graph=False, create_graph=True)
+    return gx

torchnd/modules.py

@@ -0,0 +1,635 @@
+"""N-dimensional convolution modules wrapping conv_nd functional interface."""
+
+from __future__ import annotations
+
+import math
+from abc import ABC, abstractmethod
+
+import torch
+from torch import Tensor
+from torch.utils.checkpoint import checkpoint
+
+from .functional import _normalize_tuple, adjoint_pad_nd, conv_nd, pad_nd
+
+
+class _ConvNdBase(torch.nn.Module, ABC):
+    """Base class for N-dimensional convolution modules."""
+
+    transposed: bool = False
+    in_channels: int
+    out_channels: int
+    kernel_size: tuple[int, ...]
+    dim: tuple[int, ...]
+    channel_dim: int
+    stride: tuple[int, ...]
+    padding: tuple[int, ...]
+    output_padding: tuple[int, ...]
+    dilation: tuple[int, ...]
+    groups: int
+    padding_mode: tuple[str, ...]
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int | tuple[int, ...],
+        dim: tuple[int, ...] = (-2, -1),
+        channel_dim: int = 1,
+        stride: int | tuple[int, ...] = 1,
+        padding: int | tuple[int, ...] = 0,
+        output_padding: int | tuple[int, ...] = 0,
+        dilation: int | tuple[int, ...] = 1,
+        groups: int = 1,
+        bias: bool | Tensor = True,
+        padding_mode: str | tuple[str, ...] = "zeros",
+        device: torch.device | None = None,
+        dtype: torch.dtype | None = None,
+        weight: Tensor | None = None,
+    ) -> None:
+        """Initialize N-dimensional convolution.
+
+        Parameters
+        ----------
+        in_channels
+            Number of input channels.
+        out_channels
+            Number of output channels.
+        kernel_size
+            Size of the convolving kernel.
+        dim
+            Spatial dimensions to convolve over.
+        channel_dim
+            Channel dimension index.
+        stride
+            Stride of the convolution.
+        padding
+            Padding added to each spatial dimension.
+        output_padding
+            Additional size added to output shape (transposed conv only).
+        dilation
+            Spacing between kernel elements.
+        groups
+            Number of blocked connections from input to output channels.
+        bias
+            If True, adds learnable bias. If Tensor, uses as external bias.
+        padding_mode
+            Supported padding modes are 'zeros', 'reflect', 'circular', and 'replicate'.
+            Can be a single string (applied to all dims) or a tuple of strings per dim.
+        device
+            Device for parameters. External weights/bias are moved to this device.
+        dtype
+            Data type for parameters. Complex dtypes enable complex convolution.
+            External weights/bias are cast to this dtype.
+        weight
+            If provided, uses as external weight (enables hypernetwork usage).
+            Registered as buffer to be included in state_dict.
+        """
+        super().__init__()
+
+        if in_channels % groups != 0:
+            raise ValueError(f"in_channels ({in_channels}) must be divisible by groups ({groups})")
+        if out_channels % groups != 0:
+            raise ValueError(f"out_channels ({out_channels}) must be divisible by groups ({groups})")
+
+        num_spatial = len(dim)
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = _normalize_tuple(kernel_size, num_spatial)
+        self.dim = dim
+        self.channel_dim = channel_dim
+        self.stride = _normalize_tuple(stride, num_spatial)
+        self.padding = _normalize_tuple(padding, num_spatial)
+        self.output_padding = _normalize_tuple(output_padding, num_spatial)
+        self.dilation = _normalize_tuple(dilation, num_spatial)
+        self.groups = groups
+
+        if isinstance(padding_mode, str):
+            self.padding_mode = (padding_mode,) * num_spatial
+        else:
+            self.padding_mode = tuple(padding_mode)
+
+        if self.transposed:
+            weight_shape = (in_channels, out_channels // groups, *self.kernel_size)
+        else:
+            weight_shape = (out_channels, in_channels // groups, *self.kernel_size)
+
+        for mode in self.padding_mode:
+            if mode not in ("zeros", "reflect", "circular", "replicate"):
+                raise ValueError(f"padding_mode '{mode}' is not supported")
+
+        if weight is not None:
+            weight = self._apply_device_dtype(weight, device, dtype)
+            self.register_buffer("weight", weight)
+        else:
+            self.weight = torch.nn.Parameter(torch.empty(weight_shape, device=device, dtype=dtype))
+
+        if isinstance(bias, Tensor):
+            bias = self._apply_device_dtype(bias, device, dtype)
+            self.register_buffer("bias", bias)
+        elif bias:
+            self.bias = torch.nn.Parameter(torch.empty(out_channels, device=device, dtype=dtype))
+        else:
+            self.register_buffer("bias", None)
+
+        self.reset_parameters()
+
+    @staticmethod
+    def _apply_device_dtype(
+        tensor: Tensor,
+        device: torch.device | None,
+        dtype: torch.dtype | None,
+    ) -> Tensor:
+        if device is not None or dtype is not None:
+            tensor = tensor.to(device=device, dtype=dtype)
+        return tensor
+
+    def reset_parameters(self) -> None:
+        """Initialize learnable parameters using Kaiming uniform initialization."""
+        if not isinstance(self.weight, torch.nn.Parameter):
+            return
+
+        fan_in = self.weight.shape[1] * math.prod(self.weight.shape[2:])
+
+        if self.weight.is_complex():
+            std = 1.0 / math.sqrt(fan_in * 2)
+            bound = math.sqrt(3.0) * std
+            with torch.no_grad():
+                self.weight.real.uniform_(-bound, bound)
+                self.weight.imag.uniform_(-bound, bound)
+        else:
+            torch.nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+
+        if isinstance(self.bias, torch.nn.Parameter):
+            bound = 1.0 / math.sqrt(fan_in) if fan_in > 0 else 0
+            with torch.no_grad():
+                if self.bias.is_complex():
+                    self.bias.real.uniform_(-bound, bound)
+                    self.bias.imag.uniform_(-bound, bound)
+                else:
+                    self.bias.uniform_(-bound, bound)
+
+    @abstractmethod
+    def forward(self, x: Tensor) -> Tensor:
+        """Forward pass of the convolution."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def adjoint(self, x: Tensor) -> Tensor:  # noqa: ARG002
+        """Return the adjoint convolution.
+
+        The adjoint of a convolution is the transposed convolution with
+        conjugated weights.
+
+        Parameters
+        ----------
+        x
+            Input tensor.
+        Returns
+        -------
+            Adjoint convolution.
+        """
+        raise NotImplementedError
+
+    def extra_repr(self) -> str:
+        parts = [
+            f"{self.in_channels}",
+            f"{self.out_channels}",
+            f"kernel_size={self.kernel_size}",
+            f"stride={self.stride}",
+        ]
+        if self.padding != (0,) * len(self.padding):
+            parts.append(f"padding={self.padding}")
+        if self.dilation != (1,) * len(self.dilation):
+            parts.append(f"dilation={self.dilation}")
+        if self.transposed and self.output_padding != (0,) * len(self.output_padding):
+            parts.append(f"output_padding={self.output_padding}")
+        if self.groups != 1:
+            parts.append(f"groups={self.groups}")
+        if self.bias is None:
+            parts.append("bias=False")
+        if not all(m == "zeros" for m in self.padding_mode):
+            # Show single value if all same, otherwise show tuple
+            if len(set(self.padding_mode)) == 1:
+                parts.append(f"padding_mode={self.padding_mode[0]!r}")
+            else:
+                parts.append(f"padding_mode={self.padding_mode}")
+        if self.dim != (-2, -1):
+            parts.append(f"dim={self.dim}")
+        if self.channel_dim != 1:
+            parts.append(f"channel_dim={self.channel_dim}")
+        if not isinstance(self.weight, torch.nn.Parameter):
+            parts.append("weight=external")
+        if self.bias is not None and not isinstance(self.bias, torch.nn.Parameter):
+            parts.append("bias=external")
+        return ", ".join(parts)
+
+
+class ConvNd(_ConvNdBase):
+    """N-dimensional convolution with flexible dimension specification.
+
+    Supports arbitrary spatial dimensions and dimension layouts. When native
+    PyTorch implementations (1D, 2D, 3D) are unavailable, uses recursive
+    decomposition.
+
+    Examples
+    --------
+    >>> conv = ConvNd(3, 64, kernel_size=3, dim=(-2, -1))
+    >>> x = torch.randn(2, 3, 32, 32)
+    >>> y = conv(x)
+
+    >>> # Hypernetwork usage
+    >>> external_weight = hyper_net(z)
+    >>> conv = ConvNd(3, 64, kernel_size=3, weight=external_weight, bias=False)
+    >>> y = conv(x)  # Gradients flow through hyper_net
+    """
+
+    transposed = False
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int | tuple[int, ...],
+        dim: tuple[int, ...] = (-2, -1),
+        channel_dim: int = 1,
+        stride: int | tuple[int, ...] = 1,
+        padding: int | tuple[int, ...] = 0,
+        dilation: int | tuple[int, ...] = 1,
+        groups: int = 1,
+        bias: bool | Tensor = True,
+        padding_mode: str | tuple[str, ...] = "zeros",
+        device: torch.device | None = None,
+        dtype: torch.dtype | None = None,
+        weight: Tensor | None = None,
+    ) -> None:
+        """Initialize N-dimensional convolution.
+
+        Parameters
+        ----------
+        in_channels
+            Number of input channels.
+        out_channels
+            Number of output channels.
+        kernel_size
+            Size of the convolving kernel.
+        dim
+            Spatial dimensions to convolve over.
+        channel_dim
+            Channel dimension index.
+        stride
+            Stride of the convolution.
+        padding
+            Padding added to each spatial dimension.
+        dilation
+            Spacing between kernel elements.
+        groups
+            Number of blocked connections from input to output channels.
+        bias
+            If True, adds learnable bias. If Tensor, uses as external bias.
+        padding_mode
+            Supported padding modes are 'zeros', 'reflect', 'circular', and 'replicate'.
+            Can be a single string (applied to all dims) or a tuple of strings per dim.
+        device
+            Device for parameters. External weights/bias are moved to this device.
+        dtype
+            Data type for parameters. Complex dtypes enable complex convolution.
+        weight
+            If provided, uses as external weight. Registered as buffer for state_dict.
+        """
+        super().__init__(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            dim=dim,
+            channel_dim=channel_dim,
+            stride=stride,
+            padding=padding,
+            output_padding=0,
+            dilation=dilation,
+            groups=groups,
+            bias=bias,
+            padding_mode=padding_mode,
+            device=device,
+            dtype=dtype,
+            weight=weight,
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        """Apply the forward convolution operator.
+
+        Parameters
+        ----------
+        x
+            Input tensor.
+
+        Returns
+        -------
+            Output tensor.
+        """
+        effective_padding = self.padding
+
+        if not all(m == "zeros" for m in self.padding_mode):
+            effective_padding = tuple(0 if m != "zeros" and p > 0 else p for m, p in zip(self.padding_mode, self.padding))
+            for mode in ("reflect", "replicate", "circular"):
+                dims_idx = [i for i, (m, p) in enumerate(zip(self.padding_mode, self.padding)) if m == mode and p > 0]
+                if dims_idx:
+                    spatial_dims = tuple(self.dim[i] for i in dims_idx)
+                    pad_spec = [v for i in dims_idx for v in (self.padding[i], self.padding[i])]
+                    x = checkpoint(pad_nd, x, pad_spec, mode, 0.0, spatial_dims, use_reentrant=False)
+
+        out = conv_nd(
+            x,
+            self.weight,
+            dim=self.dim,
+            channel_dim=self.channel_dim,
+            stride=self.stride,
+            padding=effective_padding,
+            dilation=self.dilation,
+            groups=self.groups,
+            transposed=False,
+        )
+
+        if self.bias is not None:
+            bias_shape = [1] * out.ndim
+            bias_shape[self.channel_dim % out.ndim] = -1
+            out = out + self.bias.view(*bias_shape)
+
+        return out
+
+    def adjoint(self, x: Tensor, input_shape: tuple[int, ...] | None = None) -> Tensor:
+        """Apply the adjoint convolution operator.
+
+        The adjoint of a convolution is the transposed convolution with
+        conjugated weights (for complex) or identical weights (for real).
+
+        This method applies the adjoint using the same weights (and does not
+        allocate a new module). Intended for iterative algorithms, not separate
+        training of an adjoint operator.
+
+        Parameters
+        ----------
+        x
+            Input tensor in the output space of the forward operator.
+        input_shape
+            Spatial shape of the input to the forward operator. Needed for
+            strided ConvNd to determine output_padding of the adjoint. For
+            stride > 1, multiple input shapes can map to the same output shape,
+            so the adjoint output size is ambiguous without this information.
+
+        Returns
+        -------
+            Result of applying the adjoint operator.
+
+        Raises
+        ------
+        ValueError
+            If bias is present (adjoint undefined with bias).
+            If input_shape is missing when any stride != 1.
+        """
+        if self.bias is not None:
+            raise ValueError("Adjoint is not defined for convolution with bias")
+
+        effective_padding = self.padding
+        if not all(m == "zeros" for m in self.padding_mode):
+            effective_padding = tuple(0 if m != "zeros" and p > 0 else p for m, p in zip(self.padding_mode, self.padding))
+
+        num_spatial = len(self.dim)
+        if input_shape is None:
+            if any(s != 1 for s in self.stride):
+                raise ValueError("input_shape is required for ConvNd.adjoint when any stride != 1")
+            out_pad = (0,) * num_spatial
+        else:
+            if len(input_shape) != num_spatial:
+                raise ValueError(f"input_shape must have length {num_spatial}, got {len(input_shape)}")
+
+            pads: list[int] = []
+            for i, (n_in, k, s, p, d) in enumerate(zip(input_shape, self.kernel_size, self.stride, self.padding, self.dilation)):
+                eff_p = effective_padding[i]
+                n_out = (n_in + 2 * p - d * (k - 1) - 1) // s + 1
+                base = (n_out - 1) * s - 2 * eff_p + d * (k - 1) + 1
+                op = (n_in + 2 * (p - eff_p)) - base
+                if op < 0 or op >= s:
+                    raise RuntimeError(
+                        "Computed output_padding is invalid. "
+                        f"Got output_padding={op} for stride={s} "
+                        f"(input={n_in}, kernel={k}, pad={p}, dilation={d})."
+                    )
+                pads.append(op)
+            out_pad = tuple(pads)
+
+        w = self.weight.conj() if self.weight.is_complex() else self.weight
+        out = conv_nd(
+            x,
+            w,
+            dim=self.dim,
+            channel_dim=self.channel_dim,
+            stride=self.stride,
+            padding=effective_padding,
+            dilation=self.dilation,
+            groups=self.groups,
+            transposed=True,
+            output_padding=out_pad,
+        )
+
+        if not all(m == "zeros" for m in self.padding_mode):
+            for mode in ("reflect", "replicate", "circular"):
+                dims_idx = [i for i, (m, p) in enumerate(zip(self.padding_mode, self.padding)) if m == mode and p > 0]
+                if dims_idx:
+                    spatial_dims = tuple(self.dim[i] for i in dims_idx)
+                    pad_spec = [v for i in dims_idx for v in (self.padding[i], self.padding[i])]
+                    original_sizes = [input_shape[i] for i in dims_idx] if input_shape else None
+                    out = adjoint_pad_nd(out, pad_spec, mode=mode, dims=spatial_dims, original_sizes=original_sizes)
+
+        return out
+
+
+class ConvTransposeNd(_ConvNdBase):
+    """N-dimensional transposed convolution with flexible dimension specification.
+
+    Also known as fractionally-strided convolution or deconvolution.
+
+    Examples
+    --------
+    >>> conv_t = ConvTransposeNd(64, 3, kernel_size=3, stride=2, dim=(-2, -1))
+    >>> x = torch.randn(2, 64, 16, 16)
+    >>> y = conv_t(x)
+    """
+
+    transposed = True
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int | tuple[int, ...],
+        dim: tuple[int, ...] = (-2, -1),
+        channel_dim: int = 1,
+        stride: int | tuple[int, ...] = 1,
+        padding: int | tuple[int, ...] = 0,
+        output_padding: int | tuple[int, ...] = 0,
+        dilation: int | tuple[int, ...] = 1,
+        groups: int = 1,
+        bias: bool | Tensor = True,
+        padding_mode: str | tuple[str, ...] = "zeros",
+        device: torch.device | None = None,
+        dtype: torch.dtype | None = None,
+        weight: Tensor | None = None,
+    ) -> None:
+        """Initialize N-dimensional transposed convolution.
+
+        Parameters
+        ----------
+        in_channels
+            Number of input channels.
+        out_channels
+            Number of output channels.
+        kernel_size
+            Size of the convolving kernel.
+        dim
+            Spatial dimensions to convolve over.
+        channel_dim
+            Channel dimension index.
+        stride
+            Stride of the convolution.
+        padding
+            Padding added to each spatial dimension.
+        output_padding
+            Additional size added to output shape to disambiguate when stride > 1.
+        dilation
+            Spacing between kernel elements.
+        groups
+            Number of blocked connections from input to output channels.
+        bias
+            If True, adds learnable bias. If Tensor, uses as external bias.
+        padding_mode
+            Supported padding modes are 'zeros', 'reflect', 'circular', and 'replicate'.
+            Can be a single string (applied to all dims) or a tuple of strings per dim.
+        device
+            Device for parameters. External weights/bias are moved to this device.
+        dtype
+            Data type for parameters. Complex dtypes enable complex convolution.
+        weight
+            If provided, uses as external weight. Registered as buffer for state_dict.
+        """
+        super().__init__(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            dim=dim,
+            channel_dim=channel_dim,
+            stride=stride,
+            padding=padding,
+            output_padding=output_padding,
+            dilation=dilation,
+            groups=groups,
+            bias=bias,
+            padding_mode=padding_mode,
+            device=device,
+            dtype=dtype,
+            weight=weight,
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        """Apply the forward convolution operator.
+
+        Parameters
+        ----------
+        x
+            Input tensor.
+
+        Returns
+        -------
+            Output tensor.
+        """
+
+        effective_padding = self.padding
+        if not all(m == "zeros" for m in self.padding_mode):
+            effective_padding = tuple(0 if m != "zeros" and p > 0 else p for m, p in zip(self.padding_mode, self.padding))
+
+        out = conv_nd(
+            x,
+            self.weight,
+            dim=self.dim,
+            channel_dim=self.channel_dim,
+            stride=self.stride,
+            padding=effective_padding,
+            dilation=self.dilation,
+            groups=self.groups,
+            transposed=True,
+            output_padding=self.output_padding,
+        )
+
+        if not all(m == "zeros" for m in self.padding_mode):
+            for mode in ("reflect", "replicate", "circular"):
+                dims_idx = [i for i, (m, p) in enumerate(zip(self.padding_mode, self.padding)) if m == mode and p > 0]
+                if dims_idx:
+                    spatial_dims = tuple(self.dim[i] for i in dims_idx)
+                    pad_spec = [v for i in dims_idx for v in (self.padding[i], self.padding[i])]
+                    original_sizes = [out.shape[self.dim[i]] - 2 * self.padding[i] for i in dims_idx]
+                    out = checkpoint(
+                        adjoint_pad_nd,
+                        out,
+                        pad_spec,
+                        mode,
+                        0.0,
+                        spatial_dims,
+                        original_sizes,
+                        use_reentrant=False,
+                    )
+
+        if self.bias is not None:
+            bias_shape = [1] * out.ndim
+            bias_shape[self.channel_dim % out.ndim] = -1
+            out = out + self.bias.view(*bias_shape)
+
+        return out
+
+    def adjoint(self, x: Tensor) -> Tensor:
+        """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).
+
+        This method applies the adjoint using the same weights (and does not
+        allocate a new module). Intended for iterative algorithms, not separate
+        training of an adjoint operator.
+
+        Parameters
+        ----------
+        x
+            Input tensor in the output space of the forward operator.
+
+        Returns
+        -------
+            Result of applying the adjoint operator.
+
+        Raises
+        ------
+        ValueError
+            If bias is present (adjoint undefined with bias).
+        """
+        if self.bias is not None:
+            raise ValueError("Adjoint is not defined for convolution with bias")
+
+        effective_padding = self.padding
+        if not all(m == "zeros" for m in self.padding_mode):
+            effective_padding = tuple(0 if m != "zeros" and p > 0 else p for m, p in zip(self.padding_mode, self.padding))
+            for mode in ("reflect", "replicate", "circular"):
+                dims_idx = [i for i, (m, p) in enumerate(zip(self.padding_mode, self.padding)) if m == mode and p > 0]
+                if dims_idx:
+                    spatial_dims = tuple(self.dim[i] for i in dims_idx)
+                    pad_spec = [v for i in dims_idx for v in (self.padding[i], self.padding[i])]
+                    x = pad_nd(x, pad_spec, mode=mode, dims=spatial_dims)
+
+        w = self.weight.conj() if self.weight.is_complex() else self.weight
+        return conv_nd(
+            x,
+            w,
+            dim=self.dim,
+            channel_dim=self.channel_dim,
+            stride=self.stride,
+            padding=effective_padding,
+            dilation=self.dilation,
+            groups=self.groups,
+            transposed=False,
+        )

torchnd-0.1.0.dist-info/METADATA

@@ -0,0 +1,154 @@
+Metadata-Version: 2.4
+Name: torchnd
+Version: 0.1.0
+Summary: ND Convolution for PyTorch
+Author-email: Felix Zimmermann <fzimmermann89@gmail.com>
+License: GPL-3.0-or-later
+Project-URL: Homepage, https://github.com/fzimmermann89/torch_conv_nd
+Project-URL: Repository, https://github.com/fzimmermann89/torch_conv_nd
+Project-URL: Issues, https://github.com/fzimmermann89/torch_conv_nd/issues
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+)
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: torch>=2.0.0
+Requires-Dist: einops>=0.6.0
+Provides-Extra: dev
+Requires-Dist: pytest>=7.4.0; extra == "dev"
+Requires-Dist: pytest-cov>=4.1.0; extra == "dev"
+Requires-Dist: ruff>=0.1.0; extra == "dev"
+Requires-Dist: mypy>=1.5.0; extra == "dev"
+Requires-Dist: pre-commit>=3.4.0; extra == "dev"
+Requires-Dist: types-setuptools; extra == "dev"
+
+# torchnd
+
+N-dimensional convolution for PyTorch that works with any number of spatial dimensions, supporting groups, dilation, transposed convolution, complex numbers, and arbitrary dimension layouts.
+
+## Installation
+
+```bash
+pip install torchnd
+```
+
+## Usage
+
+**2D convolution:**
+```python
+import torch
+from torchnd import conv_nd
+
+x = torch.randn(2, 4, 16, 16)
+weight = torch.randn(8, 4, 3, 3)
+out = conv_nd(x, weight, dim=(-2, -1), padding=1, stride=2)
+```
+
+**4D convolution** (batch × time × height × width):
+```python
+x = torch.randn(2, 4, 10, 16, 16)
+weight = torch.randn(8, 4, 3, 3, 3)
+out = conv_nd(x, weight, dim=(-3, -2, -1), padding=1)
+```
+
+**Channel-last layout:**
+```python
+x = torch.randn(2, 16, 16, 4)
+weight = torch.randn(8, 4, 3, 3)
+out = conv_nd(x, weight, dim=(1, 2), channel_dim=-1, padding=1)
+```
+
+**Transposed convolution:**
+```python
+x = torch.randn(2, 4, 8, 8)
+weight = torch.randn(4, 8, 3, 3)
+out = conv_nd(x, weight, dim=(-2, -1), padding=1, stride=2, transposed=True)
+```
+
+**Complex numbers:**
+```python
+x = torch.randn(2, 4, 16, 16, dtype=torch.complex64)
+weight = torch.randn(8, 4, 3, 3, dtype=torch.complex64)
+out = conv_nd(x, weight, dim=(-2, -1), padding=1)
+```
+
+**Asymmetric parameters per dimension:**
+```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)
+)
+```
+
+**Modules:**
+```python
+from torchnd import ConvNd, ConvTransposeNd
+
+conv = ConvNd(4, 8, 3, dim=(-2, -1), padding=1)
+out = conv(x)
+```
+
+**Padding:**
+```python
+from torchnd import pad_nd, adjoint_pad_nd
+
+# Pad or crop arbitrary dimensions
+padded = pad_nd(x, pad=(1, 1, 2, 2), dims=(-2, -1), mode="reflect")
+cropped = pad_nd(x, pad=(-1, -1), dims=(0,))  # Negative values crop
+
+# Adjoint of padding (unpads)
+unpadded = adjoint_pad_nd(padded, pad=(1, 1, 2, 2), dims=(-2, -1))
+```
+
+**Adjoint operators:**
+```python
+from torchnd import ConvNd
+
+conv = ConvNd(4, 8, 3, dim=(-2, -1), padding=1, bias=False)
+x = torch.randn(2, 4, 16, 16)
+y = torch.randn_like(conv(x))
+
+# Adjoint satisfies: <conv(x), y> = <x, conv.adjoint(y)>
+# For ConvNd: adjoint is transposed convolution with conjugated weights
+adj_y = conv.adjoint(y, input_shape=(16, 16))
+```
+
+## Functions
+
+### 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.
+
+Modes: `constant`, `reflect`, `replicate`, `circular`. Each dimension can use a different mode. For constant mode, negative padding crops symmetrically. For non-constant modes, negative padding is not supported.
+
+The `dims` parameter specifies which dimensions to pad. If `None`, pads the last N dimensions where N = len(pad) // 2. Padding is specified as pairs (left, right) per dimension.
+
+### adjoint_pad_nd
+
+Computes the adjoint of `pad_nd` via autograd. For a padding operator P, the adjoint P* satisfies the inner product identity: `<Px, y> = <x, P*y>` for all x, y. The adjoint unpads the tensor, summing contributions from padded regions back into the original shape. For constant/zeros mode, the adjoint is equivalent to cropping. For other modes, it correctly handles the adjoint of the boundary extension.
+
+### Adjoint convolution
+
+The `ConvNd` and `ConvTransposeNd` modules provide `adjoint()` methods. For `ConvNd`, the adjoint is the transposed convolution with conjugated weights (for complex) or identical weights (for real). For `ConvTransposeNd`, the adjoint is the forward convolution with conjugated weights.
+
+The adjoint satisfies the inner product identity: `<Ax, y> = <x, A*y>` where A is the convolution operator and A* is its adjoint. This is verified via dot product tests. The adjoint is undefined when bias is present.
+
+## Implementation
+
+For dimensions beyond 3D, `conv_nd` recursively decomposes the convolution into lower-dimensional operations. A 4D convolution becomes a sum of 3D convolutions applied to strided slices of the input.
+
+Complex convolution is handled by decomposing into real operations: `(a+bi)*(c+di) = (ac-bd) + (ad+bc)i`.
+
+The implementation uses native PyTorch operations when possible (1D, 2D, 3D) and falls back to recursion only for higher dimensions, minimizing memory overhead with strided views.

torchnd-0.1.0.dist-info/RECORD

@@ -0,0 +1,7 @@
+torchnd/__init__.py,sha256=H3h4vZsjgSbFKvS0w9O5kXEFVs3zqsX8eM7GdIqRmiw,308
+torchnd/functional.py,sha256=9tKawtVRUa3KFyve6pcjv603yeaVB831eXg98xio4uQ,15043
+torchnd/modules.py,sha256=wyqD55PJLEH4LUvlB4Jp3EBH14nvMsqIxvn19TB28KM,23324
+torchnd-0.1.0.dist-info/METADATA,sha256=DoVe_hcmKauvJnRBtLXQ-PoLHHCUsQoXezpzG-1k1Ko,5694
+torchnd-0.1.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+torchnd-0.1.0.dist-info/top_level.txt,sha256=ZsJixSsEHaRPdRuRlYQcnFKvii3QUbRrtEvyK_F7uIg,8
+torchnd-0.1.0.dist-info/RECORD,,

torchnd-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

torchnd-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+torchnd