[ADD] Minimal linear operator interface for PyTorch

curvlinops/_torch_base.py

@@ -0,0 +1,346 @@
+"""Defines a minimal ``LinearOperator`` interface in PyTorch."""
+
+from __future__ import annotations
+
+from typing import Callable, List, Optional, Tuple, Union
+
+import numpy
+from scipy.sparse.linalg import LinearOperator
+from torch import Size, Tensor, cat, device, dtype, from_numpy
+
+
+class PyTorchLinearOperator:
+ """Interface for linear operators in PyTorch.
+
+ Heavily inspired by the Scipy interface
+ (https://github.com/scipy/scipy/blob/v1.13.1/scipy/sparse/linalg/_interface.py),
+ but only supports a sub-set of the functionality.
+
+ One main difference is that the linear operators cannot only multiply
+ vectors/matrices specified as single PyTorch tensors, but also
+ vectors/matrices specified in tensor list format. This is common in
+ PyTorch, where the space a linear operator acts on is a tensor product
+
+ Functions that need to be implemented are ``_matmat`` and ``_adjoint``.
+
+ The interface also supports exporting the PyTorch linear operator to a SciPy linear
+ operator, which can be useful for interfacing with SciPy routines. To achieve this,
+ the functions ``_infer_device`` and ``_infer_dtype`` must be implemented.
+
+ """
+
+ def __init__(
+ self, in_shape: List[Tuple[int, ...]], out_shape: List[Tuple[int, ...]]
+ ):
+ """Store the linear operator's input and output space dimensions.
+
+ Args:
+ in_shape: A list of shapes specifying the linear operator's input space.
+ out_shape: A list of shapes specifying the linear operator's output space.
+ """
+ self._in_shape = [Size(s) for s in in_shape]
+ self._out_shape = [Size(s) for s in out_shape]
+
+ self._in_shape_flat = [s.numel() for s in self._in_shape]
+ self._out_shape_flat = [s.numel() for s in self._out_shape]
+ self.shape = (sum(self._out_shape_flat), sum(self._in_shape_flat))
+
+ def __matmul__(self, X: Union[List[Tensor], Tensor]) -> Union[List[Tensor], Tensor]:
+ """Multiply onto a vector or matrix given as PyTorch tensor or tensor list.
+
+ Args:
+ X: A vector or matrix to multiply onto, represented as a single tensor or a
+ tensor list.
+
+ Assume the linear operator has total shape ``[M, N]``:
+ If ``X`` is a single tensor, it can be of shape ``[N, K]`` (matrix), or
+ ``[N]`` (vector). The result will have shape ``[M, K]`` or ``[M]``.
+
+ Instead, we can also pass ``X`` as tensor list:
+ Assume the linear operator's rows are formed by a list of shapes
+ ``[M1, M2, ...]`` and the columns by ``[N1, N2, ...]``, such that
+ ``M1.numel() + M2.numel() + ... = M`` and ``N1.numel() + N2.numel() +
+ ... = N``. Then, ``X`` can also be a list of tensors with shape
+ ``[*N1], [*N2], ...`` (vector) or ``[*N1, K], [*N2, K], ...`` (matrix).
+ In this case, the output will be tensor list with shapes ``[*M1], [*M2],
+ ...`` (vector) or ``[K, *M1], [K, *M2], ...`` (matrix).
+
+ Returns:
+ The result of the matrix-vector or matrix-matrix multiplication in the same
+ format as ``X``.
+ """
+ # convert to tensor list format
+ X, list_format, is_vec, num_vecs = self._check_input_and_preprocess(X)
+
+ # matrix-matrix-multiply using tensor list format
+ AX = self._matmat(X)
+
+ # return same format as ``X`` passed by the user
+ return self._check_output_and_postprocess(AX, list_format, is_vec, num_vecs)
+
+ def _matmat(self, X: List[Tensor]) -> List[Tensor]:
+ """Matrix-matrix multiplication.
+
+ Args:
+ X: A list of tensors representing the matrix to multiply onto.
+ The list must contain tensors of shape ``[*N1, K], [*N2, K], ...``,
+ where ``N1, N2, ...`` are the shapes of the linear operator's columns.
+
+ Returns: # noqa: D402
+ A list of tensors with shape ``[*M1, K], [*M2, K], ...``, where ``M1, M2,
+ ...`` are the shapes of the linear operator's rows.
+
+ Raises:
+ NotImplementedError: Must be implemented by the subclass.
+ """
+ raise NotImplementedError
+
+ def adjoint(self) -> PyTorchLinearOperator:
+ """Return the adjoint of the linear operator.
+
+ Returns:
+ The adjoint of the linear operator.
+ """
+ return self._adjoint()
+
+ def _adjoint(self) -> PyTorchLinearOperator:
+ """Adjoint of the linear operator.
+
+ Returns: # noqa: D402
+ The adjoint of the linear operator.
+
+ Raises:
+ NotImplementedError: Must be implemented by the subclass.
+ """
+ raise NotImplementedError
+
+ def _check_input_and_preprocess(
+ self, X: Union[List[Tensor], Tensor]
+ ) -> Tuple[List[Tensor], bool, bool, int]:
+ """Check input format and pre-process it to a matrix in tensor list format.
+
+ Args:
+ X: The object onto which the linear operator is multiplied.
+
+ Returns:
+ X_tensor_list: The input object in tensor list format.
+ list_format: Whether the input was specified in tensor list format.
+ This is useful for post-processing the multiplication's result.
+ is_vec: Whether the input is a vector or a matrix.
+ num_vecs: The number of vectors represented by the input.
+ """
+ if isinstance(X, Tensor):
+ list_format = False
+ X_tensor_list, is_vec, num_vecs = self.__check_tensor_and_preprocess(X)
+
+ elif isinstance(X, list) and all(isinstance(x, Tensor) for x in X):
+ list_format = True
+ X_tensor_list, is_vec, num_vecs = self.__check_tensor_list_and_preprocess(X)
+
+ else:
+ raise ValueError(f"Input must be tensor or list of tensors. Got {type(X)}.")
+
+ return X_tensor_list, list_format, is_vec, num_vecs
+
+ def __check_tensor_and_preprocess(
+ self, X: Tensor
+ ) -> Tuple[List[Tensor], bool, int]:
+ """Check single-tensor input format and process into a matrix tensor list.
+
+ Args:
+ X: The tensor onto which the linear operator is multiplied.
+
+ Returns:
+ X_processed: The input tensor as matrix in tensor list format.
+ is_vec: Whether the input is a vector or a matrix.
+ num_vecs: The number of vectors represented by the input.
+
+ Raises:
+ ValueError: If the input tensor has an invalid shape.
+ """
+ if X.ndim > 2 or X.shape[0] != self.shape[1]:
+ raise ValueError(
+ f"Input tensor must have shape ({self.shape[1]},) or "
+ + f"({self.shape[1]}, K), with K arbitrary. Got {X.shape}."
+ )
+
+ # determine whether the input is a vector or matrix
+ is_vec = X.ndim == 1
+ num_vecs = 1 if is_vec else X.shape[1]
+
+ # convert to matrix in tensor list format
+ X_processed = [
+ x.reshape(*s, num_vecs)
+ for x, s in zip(X.split(self._in_shape_flat), self._in_shape)
+ ]
+
+ return X_processed, is_vec, num_vecs
+
+ def __check_tensor_list_and_preprocess(
+ self, X: List[Tensor]
+ ) -> Tuple[List[Tensor], bool, int]:
+ """Check tensor list input format and process into a matrix tensor list.
+
+ Args:
+ X: The tensor list onto which the linear operator is multiplied.
+
+ Returns:
+ X_processed: The input as matrix in tensor list format.
+ is_vec: Whether the input is a vector or a matrix.
+ num_vecs: The number of vectors represented by the input.
+
+ Raises:
+ ValueError: If the tensor entries in the list have invalid shapes.
+ """
+ if len(X) != len(self._in_shape):
+ raise ValueError(
+ f"List must contain {len(self._in_shape)} tensors. Got {len(X)}."
+ )
+
+ # check if input is a vector or a matrix
+ if all(x.shape == s for x, s in zip(X, self._in_shape)):
+ is_vec, num_vecs = True, 1
+ elif (
+ all(
+ x.ndim == len(s) + 1 and x.shape[:-1] == s
+ for x, s in zip(X, self._in_shape)
+ )
+ and len({x.shape[-1] for x in X}) == 1
+ ):
+ is_vec, (num_vecs,) = False, {x.shape[-1] for x in X}
+ else:
+ raise ValueError(
+ f"Input list must contain tensors with shapes {self._in_shape} "
+ + "and optional trailing dimension for the matrix columns. "
+ + f"Got {[x.shape for x in X]}."
+ )
+
+ # convert to matrix in tensor list format
+ X_processed = [x.unsqueeze(-1) for x in X] if is_vec else X
+
+ return X_processed, is_vec, num_vecs
+
+ def _check_output_and_postprocess(
+ self, AX: List[Tensor], list_format: bool, is_vec: bool, num_vecs: int
+ ) -> Union[List[Tensor], Tensor]:
+ """Check multiplication output and post-process it to the original format.
+
+ Args:
+ AX: The output of the multiplication as matrix in tensor list format.
+ list_format: Whether the output should be in tensor list format.
+ is_vec: Whether the output should be a vector or a matrix.
+ num_vecs: The number of vectors represented by the output.
+
+ Returns:
+ AX_processed: The output in the original format, either as single tensor
+ or list of tensors.
+
+ Raises:
+ ValueError: If the output tensor list has an invalid length or shape.
+ """
+ # verify output tensor list format
+ if len(AX) != len(self._out_shape):
+ raise ValueError(
+ f"Output list must contain {len(self._out_shape)} tensors. Got {len(AX)}."
+ )
+ if any(Ax.shape != (*s, num_vecs) for Ax, s in zip(AX, self._out_shape)):
+ raise ValueError(
+ f"Output tensors must have shapes {self._out_shape} and additional "
+ + f"trailing dimension of {num_vecs}. "
+ + f"Got {[Ax.shape for Ax in AX]}."
+ )
+
+ if list_format:
+ AX_processed = [Ax.squeeze(-1) for Ax in AX] if is_vec else AX
+ else:
+ AX_processed = cat(
+ [Ax.reshape(s, num_vecs) for Ax, s in zip(AX, self._out_shape_flat)]
+ )
+ AX_processed = AX_processed.squeeze(-1) if is_vec else AX_processed
+
+ return AX_processed
+
+ ###############################################################################
+ # SCIPY EXPORT #
+ ###############################################################################
+
+ def to_scipy(self, dtype: Optional[numpy.dtype] = None) -> LinearOperator:
+ """Wrap the PyTorch linear operator with a SciPy linear operator.
+
+ Args:
+ dtype: The data type of the SciPy linear operator. If ``None``, uses
+ NumPy's default data dtype.
+
+
+ Returns:
+ A SciPy linear operator that carries out the matrix-vector products
+ in PyTorch.
+ """
+ dev = self._infer_device()
+ dt = self._infer_dtype()
+
+ scipy_matmat = self._scipy_compatible(self.__matmul__, dev, dt)
+ A_adjoint = self.adjoint()
+ scipy_rmatmat = A_adjoint._scipy_compatible(A_adjoint.__matmul__, dev, dt)
+
+ return LinearOperator(
+ self.shape,
+ matvec=scipy_matmat,
+ rmatvec=scipy_rmatmat,
+ matmat=scipy_matmat,
+ rmatmat=scipy_rmatmat,
+ dtype=numpy.dtype(dtype) if dtype is None else dtype,
+ )
+
+ def _infer_device(self) -> device:
+ """Infer the linear operator's device.
+
+ Returns:
+ The device of the linear operator.
+
+ Raises: # noqa: D402
+ NotImplementedError: Must be implemented by subclasses.
+ """
+ raise NotImplementedError
+
+ def _infer_dtype(self) -> dtype:
+ """Infer the linear operator's data type.
+
+ Returns:
+ The data type of the linear operator.
+
+ Raises: # noqa: D402
+ NotImplementedError: Must be implemented by subclasses.
+ """
+ raise NotImplementedError
+
+ @staticmethod
+ def _scipy_compatible(
+ f: Callable[[Tensor], Tensor], device: device, dtype: dtype
+ ) -> Callable[[numpy.ndarray], numpy.ndarray]:
+ """Wrap a PyTorch matrix multiplication function to be compatible with SciPy.
+
+ Args:
+ f: The PyTorch matrix multiplication function.
+ device: The device on which the PyTorch linear operator is defined.
+ dtype: The data type of the PyTorch linear operator.
+
+ Returns:
+ A function that takes a NumPy array and returns a NumPy array.
+ """
+
+ def f_scipy(X: numpy.ndarray) -> numpy.ndarray:
+ """Scipy-compatible matrix multiplication function.
+
+ Args:
+ X: The input matrix in NumPy format.
+
+ Returns:
+ The output matrix in NumPy format.
+ """
+ X_dtype = X.dtype
+ X_torch = from_numpy(X).to(device, dtype)
+ AX_torch = f(X_torch)
+ return AX_torch.detach().cpu().numpy().astype(X_dtype)
+
+ return f_scipy