torch-named-linops 0.6.4

A flexible linear operator abstraction in pytorch.

torch-named-linops

A flexible linear operator abstraction implemented in PyTorch.

Documentation

$ pip install torch-named-linops

Quick Example

import torch
from torchlinops import Dense, Diagonal, Dim

# Create operators with named dimensions
W = torch.randn(3, 7)
A = Dense(W, Dim("MN"), ishape=Dim("N"), oshape=Dim("M"))

# Apply, take adjoint, compose
x = torch.randn(7)
y = A(x)             # Forward: y = W @ x
z = A.H(y)           # Adjoint: z = W^H @ y
w = A.N(x)           # Normal:  w = W^H @ W @ x

# Compose operators with @
d = torch.randn(3)
B = Diagonal(d, ioshape=Dim("M"))
C = B @ A             # Chain: C(x) = diag(d) @ W @ x

See the Getting Started guide for a full walkthrough.

Selected Feature List

• A dedicated abstraction for naming linear operator dimensions.
• A set of core linops, including:
  • Dense
  • Diagonal
  • FFT
  • ArrayToBlocks[^1] (similar to PyTorch's unfold but in 1D/2D/3D/arbitrary dimensions)
    • Useful for local patch extraction
  • Interpolate[^1] (similar to SigPy's interpolate/gridding)
    • Comes with kaiser_bessel and first-order spline kernels.
• .H and .N properties for adjoint $A^H$ and normal $A^HA$ linop creation.
• Chain and Add for composing linops together.
• Splitting a single linop across multiple GPUs.
• Full support for complex numbers. Adjoint takes the conjugate transpose.
• Full support for autograd-based automatic differentiation.

[^1]: Includes a functional interface and triton backend for 1D/2D/3D.

Other Packages

This package was heavily inspired by a few other influential packages. In no particular order:

• einops: named dimensions/naming things in general.
• sigpy: the linop abstraction and the idea of having dedicated adjoint and normal properties. Also inspired the NUFFT, Interpolate, and ArrayToBlocks/BlocksToArray operators.
• torch_linops: another linop abstraction. Geared more towards optimization.