torchnd 0.1.2__tar.gz → 0.2.0__tar.gz

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

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: torchnd
-Version: 0.1.2
-Summary: ND Convolution for PyTorch
+Version: 0.2.0
+Summary: ND operations 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
@@ -31,7 +31,9 @@ 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.
+N-dimensional operations for PyTorch with flexible dimension specification.
+
+`torchnd` provides convolution, transposed convolution, padding/cropping, linear interpolation, and corresponding adjoint operations for tensors with arbitrary dimension layouts. It supports complex tensors where applicable and uses native PyTorch 1D/2D/3D kernels whenever possible.
 
 ## Installation
 
@@ -112,6 +114,17 @@ cropped = pad_nd(x, pad=(-1, -1), dims=(0,))  # Negative values crop
 unpadded = adjoint_pad_nd(padded, pad=(1, 1, 2, 2), dims=(-2, -1))
 ```
 
+**Linear interpolation:**
+```python
+from torchnd import linear_interpolation_nd, adjoint_linear_interpolation_nd
+
+x = torch.randn(2, 4, 64, 64)
+low = linear_interpolation_nd(x, size=(32, 32), dims=(-2, -1))
+
+# Adjoint of the interpolation operator
+high_adjoint = adjoint_linear_interpolation_nd(low, input_size=(64, 64), dims=(-2, -1))
+```
+
 **Adjoint operators:**
 ```python
 from torchnd import ConvNd
@@ -149,6 +162,18 @@ Pad or crop a tensor to a target size, centering the original content. Combines
 
 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.
 
+### linear_interpolation_nd
+
+Resizes selected tensor dimensions using linear interpolation. The `dims` parameter specifies which dimensions are interpolated; if `None`, the last `len(size)` dimensions are used. For more than three interpolated dimensions, the operation is split into chunks of up to three dimensions to use PyTorch's native interpolation kernels.
+
+Complex tensors are handled by applying interpolation to real and imaginary parts separately.
+
+### adjoint_linear_interpolation_nd
+
+Computes the adjoint of `linear_interpolation_nd`. The adjoint is evaluated with PyTorch's native interpolation backward kernels and satisfies the inner product identity `<Ax, y> = <x, A*y>` for the linear interpolation operator A.
+
+The adjoint supports autograd, forward-mode AD, and `torch.func.vmap`.
+
 ### 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.
@@ -159,6 +184,8 @@ The adjoint satisfies the inner product identity: `<Ax, y> = <x, A*y>` where A i
 
 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.
 
+For interpolation of more than three dimensions, `linear_interpolation_nd` applies native PyTorch interpolation to chunks of up to three dimensions. The adjoint applies the corresponding chunks in reverse order.
+
 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.2 → torchnd-0.2.0}/README.md

@@ -1,6 +1,8 @@
 # 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.
+N-dimensional operations for PyTorch with flexible dimension specification.
+
+`torchnd` provides convolution, transposed convolution, padding/cropping, linear interpolation, and corresponding adjoint operations for tensors with arbitrary dimension layouts. It supports complex tensors where applicable and uses native PyTorch 1D/2D/3D kernels whenever possible.
 
 ## Installation
 
@@ -81,6 +83,17 @@ cropped = pad_nd(x, pad=(-1, -1), dims=(0,))  # Negative values crop
 unpadded = adjoint_pad_nd(padded, pad=(1, 1, 2, 2), dims=(-2, -1))
 ```
 
+**Linear interpolation:**
+```python
+from torchnd import linear_interpolation_nd, adjoint_linear_interpolation_nd
+
+x = torch.randn(2, 4, 64, 64)
+low = linear_interpolation_nd(x, size=(32, 32), dims=(-2, -1))
+
+# Adjoint of the interpolation operator
+high_adjoint = adjoint_linear_interpolation_nd(low, input_size=(64, 64), dims=(-2, -1))
+```
+
 **Adjoint operators:**
 ```python
 from torchnd import ConvNd
@@ -118,6 +131,18 @@ Pad or crop a tensor to a target size, centering the original content. Combines
 
 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.
 
+### linear_interpolation_nd
+
+Resizes selected tensor dimensions using linear interpolation. The `dims` parameter specifies which dimensions are interpolated; if `None`, the last `len(size)` dimensions are used. For more than three interpolated dimensions, the operation is split into chunks of up to three dimensions to use PyTorch's native interpolation kernels.
+
+Complex tensors are handled by applying interpolation to real and imaginary parts separately.
+
+### adjoint_linear_interpolation_nd
+
+Computes the adjoint of `linear_interpolation_nd`. The adjoint is evaluated with PyTorch's native interpolation backward kernels and satisfies the inner product identity `<Ax, y> = <x, A*y>` for the linear interpolation operator A.
+
+The adjoint supports autograd, forward-mode AD, and `torch.func.vmap`.
+
 ### 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.
@@ -128,6 +153,8 @@ The adjoint satisfies the inner product identity: `<Ax, y> = <x, A*y>` where A i
 
 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.
 
+For interpolation of more than three dimensions, `linear_interpolation_nd` applies native PyTorch interpolation to chunks of up to three dimensions. The adjoint applies the corresponding chunks in reverse order.
+
 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.2 → torchnd-0.2.0}/pyproject.toml

@@ -4,8 +4,8 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "torchnd"
-version = "0.1.2"
-description = "ND Convolution for PyTorch"
+version = "0.2.0"
+description = "ND operations for PyTorch"
 readme = "README.md"
 requires-python = ">=3.10"
 license = {text = "GPL-3.0-or-later"}

torchnd-0.2.0/src/torchnd/__init__.py

@@ -0,0 +1,27 @@
+__version__ = "0.2.0"
+
+from torchnd.functional import (
+    adjoint_linear_interpolation_nd,
+    adjoint_nearest_interpolation_nd,
+    adjoint_pad_nd,
+    conv_nd,
+    linear_interpolation_nd,
+    nearest_interpolation_nd,
+    pad_nd,
+    pad_or_crop_to_size,
+)
+from torchnd.modules import ConvNd, ConvTransposeNd
+
+__all__ = [
+    "adjoint_linear_interpolation_nd",
+    "adjoint_nearest_interpolation_nd",
+    "adjoint_pad_nd",
+    "conv_nd",
+    "linear_interpolation_nd",
+    "nearest_interpolation_nd",
+    "pad_nd",
+    "pad_or_crop_to_size",
+    "ConvNd",
+    "ConvTransposeNd",
+    "__version__",
+]

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

@@ -18,11 +18,318 @@ CONV_TRANSPOSE_REGISTRY: dict[int, Callable[..., Tensor]] = {
     3: torch.nn.functional.conv_transpose3d,
 }
 
+LINEAR_INTERPOLATION_MODE: dict[int, str] = {
+    1: "linear",
+    2: "bilinear",
+    3: "trilinear",
+}
+
+LINEAR_INTERPOLATION_ADJOINT_REGISTRY: dict[int, Callable[..., Tensor]] = {
+    1: torch.ops.aten.upsample_linear1d_backward,
+    2: torch.ops.aten.upsample_bilinear2d_backward,
+    3: torch.ops.aten.upsample_trilinear3d_backward,
+}
+
+NEAREST_INTERPOLATION_ADJOINT_REGISTRY: dict[int, Callable[..., Tensor]] = {
+    1: torch.ops.aten.upsample_nearest1d_backward,
+    2: torch.ops.aten.upsample_nearest2d_backward,
+    3: torch.ops.aten.upsample_nearest3d_backward,
+}
+
 
 def _normalize_tuple(value: int | tuple[int, ...], length: int) -> tuple[int, ...]:
     return (value,) * length if isinstance(value, int) else tuple(value)
 
 
+def _normalize_dims(ndim: int, dims: Sequence[int]) -> tuple[int, ...]:
+    normalized = tuple(d % ndim for d in dims)
+    if len(set(normalized)) != len(normalized):
+        raise ValueError("duplicate dimensions not supported")
+    return normalized
+
+
+def _linear_interpolation_nd_real(
+    x: Tensor,
+    size: Sequence[int],
+    dims: Sequence[int] | None = None,
+    align_corners: bool = False,
+) -> Tensor:
+    size = tuple(size)
+    if dims is None:
+        if len(size) > x.ndim:
+            raise ValueError("len(size) must be <= x.ndim")
+        target_dims = tuple(range(x.ndim - len(size), x.ndim))
+    elif len(dims) != len(size):
+        raise ValueError("len(dims) must equal len(size)")
+    else:
+        target_dims = _normalize_dims(x.ndim, dims)
+
+    for i in range(0, len(target_dims), 3):
+        chunk_dims = target_dims[i : i + 3]
+        chunk_size = size[i : i + 3]
+
+        other_dims = tuple(dim for dim in range(x.ndim) if dim not in chunk_dims)
+        to_standard = (*other_dims, *chunk_dims)
+        from_standard = tuple(sorted(range(len(to_standard)), key=to_standard.__getitem__))
+        other_shape = tuple(x.shape[dim] for dim in other_dims)
+
+        x = x.permute(to_standard)
+        x = x.flatten(end_dim=len(other_dims) - 1).unsqueeze(0) if other_dims else x.unsqueeze(0).unsqueeze(0)
+
+        x = torch.nn.functional.interpolate(
+            x,
+            size=chunk_size,
+            mode=LINEAR_INTERPOLATION_MODE[len(chunk_dims)],
+            align_corners=align_corners,
+        )
+
+        x = x.squeeze(0).unflatten(0, other_shape).permute(from_standard) if other_dims else x.squeeze(0).squeeze(0)
+    return x
+
+
+def linear_interpolation_nd(
+    x: Tensor,
+    size: Sequence[int],
+    dims: Sequence[int] | None = None,
+    align_corners: bool = False,
+) -> Tensor:
+    """Resize selected tensor dimensions using linear interpolation.
+
+    The selected dimensions are moved to PyTorch's interpolation layout and
+    processed in chunks of up to three dimensions. Complex tensors are
+    interpolated component-wise.
+
+    Parameters
+    ----------
+    x
+        Input tensor.
+    size
+        Target size for the interpolated dimensions.
+    dims
+        Dimensions to interpolate. If None, interpolates the last ``len(size)`` dimensions.
+    align_corners
+        Geometric convention passed to ``torch.nn.functional.interpolate``.
+
+    Returns
+    -------
+        Tensor with selected dimensions resized to ``size``.
+    """
+    if len(size) == 0:
+        return x
+    if x.is_complex():
+        return torch.complex(
+            _linear_interpolation_nd_real(x.real, size, dims, align_corners),
+            _linear_interpolation_nd_real(x.imag, size, dims, align_corners),
+        )
+    return _linear_interpolation_nd_real(x, size, dims, align_corners)
+
+
+def _nearest_interpolation_nd_real(
+    x: Tensor,
+    size: Sequence[int],
+    dims: Sequence[int] | None = None,
+) -> Tensor:
+    size = tuple(size)
+    if dims is None:
+        if len(size) > x.ndim:
+            raise ValueError("len(size) must be <= x.ndim")
+        target_dims = tuple(range(x.ndim - len(size), x.ndim))
+    elif len(dims) != len(size):
+        raise ValueError("len(dims) must equal len(size)")
+    else:
+        target_dims = _normalize_dims(x.ndim, dims)
+
+    for i in range(0, len(target_dims), 3):
+        chunk_dims = target_dims[i : i + 3]
+        chunk_size = size[i : i + 3]
+
+        other_dims = tuple(dim for dim in range(x.ndim) if dim not in chunk_dims)
+        to_standard = (*other_dims, *chunk_dims)
+        from_standard = tuple(sorted(range(len(to_standard)), key=to_standard.__getitem__))
+        other_shape = tuple(x.shape[dim] for dim in other_dims)
+
+        x = x.permute(to_standard)
+        x = x.flatten(end_dim=len(other_dims) - 1).unsqueeze(0) if other_dims else x.unsqueeze(0).unsqueeze(0)
+        x = torch.nn.functional.interpolate(x, size=chunk_size, mode="nearest")
+        x = x.squeeze(0).unflatten(0, other_shape).permute(from_standard) if other_dims else x.squeeze(0).squeeze(0)
+    return x
+
+
+def nearest_interpolation_nd(
+    x: Tensor,
+    size: Sequence[int],
+    dims: Sequence[int] | None = None,
+) -> Tensor:
+    """Resize selected tensor dimensions using nearest-neighbor interpolation.
+
+    Parameters
+    ----------
+    x
+        Input tensor.
+    size
+        Target size for the interpolated dimensions.
+    dims
+        Dimensions to interpolate. If None, interpolates the last ``len(size)`` dimensions.
+
+    Returns
+    -------
+        Tensor with selected dimensions resized to ``size``.
+    """
+    if len(size) == 0:
+        return x
+    if x.is_complex():
+        return torch.complex(
+            _nearest_interpolation_nd_real(x.real, size, dims),
+            _nearest_interpolation_nd_real(x.imag, size, dims),
+        )
+    return _nearest_interpolation_nd_real(x, size, dims)
+
+
+def _adjoint_linear_interpolation_nd_real(
+    x: Tensor,
+    input_size: Sequence[int],
+    dims: Sequence[int] | None = None,
+    align_corners: bool = False,
+) -> Tensor:
+    input_size = tuple(input_size)
+    if dims is None:
+        if len(input_size) > x.ndim:
+            raise ValueError("len(input_size) must be <= x.ndim")
+        target_dims = tuple(range(x.ndim - len(input_size), x.ndim))
+    elif len(dims) != len(input_size):
+        raise ValueError("len(dims) must equal len(input_size)")
+    else:
+        target_dims = _normalize_dims(x.ndim, dims)
+
+    chunks = [(target_dims[i : i + 3], input_size[i : i + 3]) for i in range(0, len(target_dims), 3)]
+    for chunk_dims, chunk_input_size in reversed(chunks):
+        output_size = tuple(x.shape[dim] for dim in chunk_dims)
+
+        other_dims = tuple(dim for dim in range(x.ndim) if dim not in chunk_dims)
+        to_standard = (*other_dims, *chunk_dims)
+        from_standard = tuple(sorted(range(len(to_standard)), key=to_standard.__getitem__))
+        other_shape = tuple(x.shape[dim] for dim in other_dims)
+
+        x = x.permute(to_standard)
+        x = x.flatten(end_dim=len(other_dims) - 1).unsqueeze(0) if other_dims else x.unsqueeze(0).unsqueeze(0)
+
+        x = LINEAR_INTERPOLATION_ADJOINT_REGISTRY[len(chunk_input_size)](
+            x,
+            list(output_size),
+            [x.shape[0], x.shape[1], *chunk_input_size],
+            align_corners,
+            *([None] * len(chunk_input_size)),
+        )
+
+        x = x.squeeze(0).unflatten(0, other_shape).permute(from_standard) if other_dims else x.squeeze(0).squeeze(0)
+    return x
+
+
+def adjoint_linear_interpolation_nd(
+    x: Tensor,
+    input_size: Sequence[int],
+    dims: Sequence[int] | None = None,
+    align_corners: bool = False,
+) -> Tensor:
+    """Apply the adjoint of ``linear_interpolation_nd``.
+
+    The adjoint is evaluated with PyTorch's native interpolation backward
+    kernels, not by running a dummy forward interpolation. Complex tensors are
+    handled component-wise.
+
+    Parameters
+    ----------
+    x
+        Tensor in the output space of ``linear_interpolation_nd``.
+    input_size
+        Size of the selected dimensions before interpolation.
+    dims
+        Dimensions that were interpolated. If None, uses the last ``len(input_size)`` dimensions.
+    align_corners
+        Geometric convention passed to ``torch.nn.functional.interpolate``.
+
+    Returns
+    -------
+        Tensor in the input space of ``linear_interpolation_nd``.
+    """
+    if len(input_size) == 0:
+        return x
+    if x.is_complex():
+        return torch.complex(
+            _adjoint_linear_interpolation_nd_real(x.real, input_size, dims, align_corners),
+            _adjoint_linear_interpolation_nd_real(x.imag, input_size, dims, align_corners),
+        )
+    return _adjoint_linear_interpolation_nd_real(x, input_size, dims, align_corners)
+
+
+def _adjoint_nearest_interpolation_nd_real(
+    x: Tensor,
+    input_size: Sequence[int],
+    dims: Sequence[int] | None = None,
+) -> Tensor:
+    input_size = tuple(input_size)
+    if dims is None:
+        if len(input_size) > x.ndim:
+            raise ValueError("len(input_size) must be <= x.ndim")
+        target_dims = tuple(range(x.ndim - len(input_size), x.ndim))
+    elif len(dims) != len(input_size):
+        raise ValueError("len(dims) must equal len(input_size)")
+    else:
+        target_dims = _normalize_dims(x.ndim, dims)
+
+    chunks = [(target_dims[i : i + 3], input_size[i : i + 3]) for i in range(0, len(target_dims), 3)]
+    for chunk_dims, chunk_input_size in reversed(chunks):
+        output_size = tuple(x.shape[dim] for dim in chunk_dims)
+
+        other_dims = tuple(dim for dim in range(x.ndim) if dim not in chunk_dims)
+        to_standard = (*other_dims, *chunk_dims)
+        from_standard = tuple(sorted(range(len(to_standard)), key=to_standard.__getitem__))
+        other_shape = tuple(x.shape[dim] for dim in other_dims)
+
+        x = x.permute(to_standard)
+        x = x.flatten(end_dim=len(other_dims) - 1).unsqueeze(0) if other_dims else x.unsqueeze(0).unsqueeze(0)
+
+        x = NEAREST_INTERPOLATION_ADJOINT_REGISTRY[len(chunk_input_size)](
+            x,
+            list(output_size),
+            [x.shape[0], x.shape[1], *chunk_input_size],
+            *([None] * len(chunk_input_size)),
+        )
+
+        x = x.squeeze(0).unflatten(0, other_shape).permute(from_standard) if other_dims else x.squeeze(0).squeeze(0)
+    return x
+
+
+def adjoint_nearest_interpolation_nd(
+    x: Tensor,
+    input_size: Sequence[int],
+    dims: Sequence[int] | None = None,
+) -> Tensor:
+    """Apply the adjoint of ``nearest_interpolation_nd``.
+
+    Parameters
+    ----------
+    x
+        Tensor in the output space of ``nearest_interpolation_nd``.
+    input_size
+        Size of the selected dimensions before interpolation.
+    dims
+        Dimensions that were interpolated. If None, uses the last ``len(input_size)`` dimensions.
+
+    Returns
+    -------
+        Tensor in the input space of ``nearest_interpolation_nd``.
+    """
+    if len(input_size) == 0:
+        return x
+    if x.is_complex():
+        return torch.complex(
+            _adjoint_nearest_interpolation_nd_real(x.real, input_size, dims),
+            _adjoint_nearest_interpolation_nd_real(x.imag, input_size, dims),
+        )
+    return _adjoint_nearest_interpolation_nd_real(x, input_size, dims)
+
+
 def conv_nd(
     x: Tensor,
     weight: Tensor,
@@ -295,10 +602,8 @@ def pad_nd(
     else:
         if len(dims) != n_dims:
             raise ValueError("len(dims) must equal len(pad) // 2")
-        target_dims = [d % ndim for d in dims]
+        target_dims = list(_normalize_dims(ndim, 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(x, target_dims, pad_pairs, value)
@@ -340,7 +645,7 @@ def pad_or_crop_to_size(
     else:
         if len(dims) != len(size):
             raise ValueError("len(dims) must equal len(size)")
-        target_dims = [d % ndim for d in dims]
+        target_dims = list(_normalize_dims(ndim, dims))
 
     pad_list: list[int] = []
     for dim, target in zip(target_dims, size, strict=True):
@@ -374,10 +679,6 @@ def _pad_non_constant(
     pad_pairs: list[tuple[int, int]],
     mode: str,
 ) -> Tensor:
-    for left, right in pad_pairs:
-        if left < 0 or right < 0:
-            raise ValueError(f"Negative padding is not supported for mode '{mode}'. Use mode='constant' for cropping.")
-
     result = x
 
     for i in range(0, len(target_dims), 3):
@@ -464,7 +765,7 @@ def adjoint_pad_nd(
     else:
         if len(dims) != n_dims:
             raise ValueError("len(dims) must equal len(pad) // 2")
-        target_dims = [d % ndim for d in dims]
+        target_dims = list(_normalize_dims(ndim, dims))
 
     if original_sizes is None:
         original_sizes = []

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

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: torchnd
-Version: 0.1.2
-Summary: ND Convolution for PyTorch
+Version: 0.2.0
+Summary: ND operations 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
@@ -31,7 +31,9 @@ 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.
+N-dimensional operations for PyTorch with flexible dimension specification.
+
+`torchnd` provides convolution, transposed convolution, padding/cropping, linear interpolation, and corresponding adjoint operations for tensors with arbitrary dimension layouts. It supports complex tensors where applicable and uses native PyTorch 1D/2D/3D kernels whenever possible.
 
 ## Installation
 
@@ -112,6 +114,17 @@ cropped = pad_nd(x, pad=(-1, -1), dims=(0,))  # Negative values crop
 unpadded = adjoint_pad_nd(padded, pad=(1, 1, 2, 2), dims=(-2, -1))
 ```
 
+**Linear interpolation:**
+```python
+from torchnd import linear_interpolation_nd, adjoint_linear_interpolation_nd
+
+x = torch.randn(2, 4, 64, 64)
+low = linear_interpolation_nd(x, size=(32, 32), dims=(-2, -1))
+
+# Adjoint of the interpolation operator
+high_adjoint = adjoint_linear_interpolation_nd(low, input_size=(64, 64), dims=(-2, -1))
+```
+
 **Adjoint operators:**
 ```python
 from torchnd import ConvNd
@@ -149,6 +162,18 @@ Pad or crop a tensor to a target size, centering the original content. Combines
 
 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.
 
+### linear_interpolation_nd
+
+Resizes selected tensor dimensions using linear interpolation. The `dims` parameter specifies which dimensions are interpolated; if `None`, the last `len(size)` dimensions are used. For more than three interpolated dimensions, the operation is split into chunks of up to three dimensions to use PyTorch's native interpolation kernels.
+
+Complex tensors are handled by applying interpolation to real and imaginary parts separately.
+
+### adjoint_linear_interpolation_nd
+
+Computes the adjoint of `linear_interpolation_nd`. The adjoint is evaluated with PyTorch's native interpolation backward kernels and satisfies the inner product identity `<Ax, y> = <x, A*y>` for the linear interpolation operator A.
+
+The adjoint supports autograd, forward-mode AD, and `torch.func.vmap`.
+
 ### 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.
@@ -159,6 +184,8 @@ The adjoint satisfies the inner product identity: `<Ax, y> = <x, A*y>` where A i
 
 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.
 
+For interpolation of more than three dimensions, `linear_interpolation_nd` applies native PyTorch interpolation to chunks of up to three dimensions. The adjoint applies the corresponding chunks in reverse order.
+
 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.2 → torchnd-0.2.0}/tests/test_functional.py

@@ -5,7 +5,16 @@ from collections.abc import Iterator
 import pytest
 import torch
 
-from torchnd import adjoint_pad_nd, conv_nd, pad_nd, pad_or_crop_to_size
+from torchnd import (
+    adjoint_linear_interpolation_nd,
+    adjoint_nearest_interpolation_nd,
+    adjoint_pad_nd,
+    conv_nd,
+    linear_interpolation_nd,
+    nearest_interpolation_nd,
+    pad_nd,
+    pad_or_crop_to_size,
+)
 from torchnd.functional import CONV_REGISTRY, CONV_TRANSPOSE_REGISTRY
 
 
@@ -373,10 +382,12 @@ class TestPadNd:
             pad_nd(torch.zeros(2, 2), **kwargs)
 
     @pytest.mark.parametrize("mode", ["reflect", "replicate", "circular"])
-    def test_negative_padding_non_constant_raises(self, mode: str) -> None:
-        """Verify negative padding with non-constant mode raises ValueError."""
-        with pytest.raises(ValueError, match="Negative padding is not supported"):
-            pad_nd(torch.zeros(5, 5), (-1, 0), dims=(0,), mode=mode)
+    def test_negative_padding_non_constant_crops(self, mode: str) -> None:
+        """Verify negative padding works for non-constant modes."""
+        x = torch.arange(5).float()[None, None, :]
+        out = pad_nd(x, (-1, 1), dims=(2,), mode=mode)
+        expected = torch.nn.functional.pad(x, (-1, 1), mode=mode)
+        torch.testing.assert_close(out, expected)
 
     def test_multi_dim_non_constant_chunking(self) -> None:
         """Verify correct handling of >3 dims requiring multiple pad calls."""
@@ -527,3 +538,219 @@ class TestAdjointPadNd:
         result = adjoint_pad_nd(padded, pad, mode="reflect", dims=dims)
 
         assert result.shape == (2, 4, 10, 10)
+
+
+class TestLinearInterpolationNd:
+    """Tests for linear_interpolation_nd and adjoint_linear_interpolation_nd."""
+
+    @pytest.mark.parametrize(
+        "shape, size, dims",
+        [
+            ((2, 3, 16), (7,), (-1,)),
+            ((2, 3, 16, 18), (7, 9), (-2, -1)),
+            ((2, 3, 8, 10, 12), (4, 5, 6), (-3, -2, -1)),
+        ],
+    )
+    def test_matches_torch_interpolate(
+        self,
+        shape: tuple[int, ...],
+        size: tuple[int, ...],
+        dims: tuple[int, ...],
+    ) -> None:
+        """Verify standard layouts match torch.nn.functional.interpolate."""
+        x = torch.randn(shape)
+        expected = torch.nn.functional.interpolate(
+            x,
+            size=size,
+            mode={1: "linear", 2: "bilinear", 3: "trilinear"}[len(size)],
+            align_corners=False,
+        )
+        result = linear_interpolation_nd(x, size, dims=dims)
+        torch.testing.assert_close(result, expected)
+
+    def test_arbitrary_dims(self) -> None:
+        """Verify interpolation works when spatial dimensions are not trailing."""
+        x = torch.randn(2, 12, 3, 14)
+        result = linear_interpolation_nd(x, size=(6, 7), dims=(1, 3))
+        assert result.shape == (2, 6, 3, 7)
+
+    def test_complex(self) -> None:
+        """Verify complex tensors are interpolated component-wise."""
+        x = torch.randn(2, 3, 12, 14, dtype=torch.complex64)
+        result = linear_interpolation_nd(x, size=(6, 7), dims=(-2, -1))
+        expected = torch.complex(
+            linear_interpolation_nd(x.real, size=(6, 7), dims=(-2, -1)),
+            linear_interpolation_nd(x.imag, size=(6, 7), dims=(-2, -1)),
+        )
+        torch.testing.assert_close(result, expected)
+
+    def test_more_than_three_dims(self) -> None:
+        """Verify interpolation with more than three dimensions uses multiple backend calls."""
+        x = torch.randn(2, 3, 4, 5, 6)
+        result = linear_interpolation_nd(x, size=(3, 4, 5, 7), dims=(1, 2, 3, 4))
+        assert result.shape == (2, 3, 4, 5, 7)
+
+    @pytest.mark.parametrize(
+        "shape, size, dims, dtype",
+        [
+            ((2, 3, 16), (7,), (-1,), torch.float32),
+            ((2, 12, 3, 14), (6, 7), (1, 3), torch.float32),
+            ((2, 3, 8, 10, 12), (4, 5, 6), (-3, -2, -1), torch.complex64),
+            ((2, 3, 4, 5, 6), (4, 3, 4, 5), (1, 2, 3, 4), torch.float32),
+        ],
+    )
+    def test_adjointness(
+        self,
+        shape: tuple[int, ...],
+        size: tuple[int, ...],
+        dims: tuple[int, ...],
+        dtype: torch.dtype,
+    ) -> None:
+        """Verify <Ax, y> == <x, A^H y>."""
+        x = torch.randn(shape, dtype=dtype)
+        y_shape = list(shape)
+        normalized_dims = tuple(dim % len(shape) for dim in dims)
+        for dim, target in zip(normalized_dims, size, strict=True):
+            y_shape[dim] = target
+        y = torch.randn(y_shape, dtype=dtype)
+
+        ax = linear_interpolation_nd(x, size, dims=dims)
+        ahy = adjoint_linear_interpolation_nd(y, [shape[dim] for dim in normalized_dims], dims=dims)
+
+        lhs = torch.vdot(ax.flatten(), y.flatten())
+        rhs = torch.vdot(x.flatten(), ahy.flatten())
+        torch.testing.assert_close(lhs, rhs, rtol=1e-4, atol=1e-5)
+
+    def test_adjoint_is_differentiable(self) -> None:
+        """Verify the adjoint operation is differentiable."""
+        y = torch.randn(2, 3, 6, 7, requires_grad=True)
+        result = adjoint_linear_interpolation_nd(y, input_size=(12, 14), dims=(-2, -1))
+        result.square().sum().backward()
+        assert y.grad is not None
+
+    def test_gradcheck(self) -> None:
+        """Verify numerical gradients for interpolation and its adjoint."""
+        x = torch.randn(1, 2, 5, 6, dtype=torch.float64, requires_grad=True)
+        y = torch.randn(1, 2, 3, 4, dtype=torch.float64, requires_grad=True)
+
+        assert torch.autograd.gradcheck(lambda value: linear_interpolation_nd(value, size=(3, 4), dims=(-2, -1)), x)
+        assert torch.autograd.gradcheck(
+            lambda value: adjoint_linear_interpolation_nd(value, input_size=(5, 6), dims=(-2, -1)),
+            y,
+        )
+
+    def test_forward_jvp(self) -> None:
+        """Verify interpolation works with forward-mode AD."""
+        x = torch.randn(2, 3, 12, 14)
+        tangent = torch.randn_like(x)
+        primal, tangent_out = torch.func.jvp(
+            lambda y: linear_interpolation_nd(y, size=(6, 7), dims=(-2, -1)),
+            (x,),
+            (tangent,),
+        )
+        assert primal.shape == (2, 3, 6, 7)
+        torch.testing.assert_close(tangent_out, linear_interpolation_nd(tangent, size=(6, 7), dims=(-2, -1)))
+
+    def test_forward_vmap(self) -> None:
+        """Verify interpolation works with vmap."""
+        x = torch.randn(5, 2, 3, 12, 14)
+        result = torch.func.vmap(lambda y: linear_interpolation_nd(y, size=(6, 7), dims=(-2, -1)))(x)
+        assert result.shape == (5, 2, 3, 6, 7)
+
+    def test_adjoint_jvp(self) -> None:
+        """Verify the adjoint works with forward-mode AD."""
+        y = torch.randn(2, 3, 6, 7)
+        tangent = torch.randn_like(y)
+        primal, tangent_out = torch.func.jvp(
+            lambda x: adjoint_linear_interpolation_nd(x, input_size=(12, 14), dims=(-2, -1)),
+            (y,),
+            (tangent,),
+        )
+        assert primal.shape == (2, 3, 12, 14)
+        torch.testing.assert_close(
+            tangent_out,
+            adjoint_linear_interpolation_nd(tangent, input_size=(12, 14), dims=(-2, -1)),
+        )
+
+    def test_adjoint_vmap(self) -> None:
+        """Verify the adjoint works with vmap."""
+        y = torch.randn(5, 2, 3, 6, 7)
+        result = torch.func.vmap(lambda x: adjoint_linear_interpolation_nd(x, input_size=(12, 14), dims=(-2, -1)))(y)
+        assert result.shape == (5, 2, 3, 12, 14)
+
+
+class TestNearestInterpolationNd:
+    """Tests for nearest_interpolation_nd and adjoint_nearest_interpolation_nd."""
+
+    @pytest.mark.parametrize(
+        "shape, size, dims",
+        [
+            ((2, 3, 16), (7,), (-1,)),
+            ((2, 3, 16, 18), (7, 9), (-2, -1)),
+            ((2, 3, 8, 10, 12), (4, 5, 6), (-3, -2, -1)),
+        ],
+    )
+    def test_matches_torch_interpolate(
+        self,
+        shape: tuple[int, ...],
+        size: tuple[int, ...],
+        dims: tuple[int, ...],
+    ) -> None:
+        """Verify standard layouts match torch.nn.functional.interpolate."""
+        x = torch.randn(shape)
+        expected = torch.nn.functional.interpolate(x, size=size, mode="nearest")
+        result = nearest_interpolation_nd(x, size, dims=dims)
+        torch.testing.assert_close(result, expected)
+
+    def test_arbitrary_dims(self) -> None:
+        """Verify interpolation works when spatial dimensions are not trailing."""
+        x = torch.randn(2, 12, 3, 14)
+        result = nearest_interpolation_nd(x, size=(6, 7), dims=(1, 3))
+        assert result.shape == (2, 6, 3, 7)
+
+    def test_complex(self) -> None:
+        """Verify complex tensors are interpolated component-wise."""
+        x = torch.randn(2, 3, 12, 14, dtype=torch.complex64)
+        result = nearest_interpolation_nd(x, size=(6, 7), dims=(-2, -1))
+        expected = torch.complex(
+            nearest_interpolation_nd(x.real, size=(6, 7), dims=(-2, -1)),
+            nearest_interpolation_nd(x.imag, size=(6, 7), dims=(-2, -1)),
+        )
+        torch.testing.assert_close(result, expected)
+
+    def test_more_than_three_dims(self) -> None:
+        """Verify interpolation with more than three dimensions uses multiple backend calls."""
+        x = torch.randn(2, 3, 4, 5, 6)
+        result = nearest_interpolation_nd(x, size=(3, 4, 5, 7), dims=(1, 2, 3, 4))
+        assert result.shape == (2, 3, 4, 5, 7)
+
+    @pytest.mark.parametrize(
+        "shape, size, dims, dtype",
+        [
+            ((2, 3, 16), (7,), (-1,), torch.float32),
+            ((2, 12, 3, 14), (6, 7), (1, 3), torch.float32),
+            ((2, 3, 8, 10, 12), (4, 5, 6), (-3, -2, -1), torch.complex64),
+            ((2, 3, 4, 5, 6), (4, 3, 4, 5), (1, 2, 3, 4), torch.float32),
+        ],
+    )
+    def test_adjointness(
+        self,
+        shape: tuple[int, ...],
+        size: tuple[int, ...],
+        dims: tuple[int, ...],
+        dtype: torch.dtype,
+    ) -> None:
+        """Verify <Ax, y> == <x, A^H y>."""
+        x = torch.randn(shape, dtype=dtype)
+        y_shape = list(shape)
+        normalized_dims = tuple(dim % len(shape) for dim in dims)
+        for dim, target in zip(normalized_dims, size, strict=True):
+            y_shape[dim] = target
+        y = torch.randn(y_shape, dtype=dtype)
+
+        ax = nearest_interpolation_nd(x, size, dims=dims)
+        ahy = adjoint_nearest_interpolation_nd(y, [shape[dim] for dim in normalized_dims], dims=dims)
+
+        lhs = torch.vdot(ax.flatten(), y.flatten())
+        rhs = torch.vdot(x.flatten(), ahy.flatten())
+        torch.testing.assert_close(lhs, rhs, rtol=1e-4, atol=1e-5)

torchnd-0.1.2/src/torchnd/__init__.py

@@ -1,14 +0,0 @@
-__version__ = "0.1.2"
-
-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__",
-]