torch-named-linops 0.6.5

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¹ (similar to PyTorch's unfold but in 1D/2D/3D/arbitrary dimensions)
    • Useful for local patch extraction
  • Interpolate¹ (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.

Documentation

The documentation is built using Zensical with tutorials authored as Marimo notebooks.

Building Documentation Locally

Install just (a command runner):

# Install with uv
uv tool install just

# Or with other package managers
# Homebrew: brew install just
# Cargo: cargo install just

Build the documentation:

# Build tutorials and documentation
just docs

# Or serve locally with auto-reload
just dev

Available commands:

Command	Description
just tutorials	Build tutorials from marimo notebooks
just docs	Build full documentation (includes tutorials)
just serve	Serve documentation locally
just dev	Build and serve documentation

Editing Tutorials

Tutorials are marimo notebooks in tutorials/*.py. To edit:

# Edit a tutorial interactively
uv run marimo edit tutorials/basics.py

# Rebuild markdown after edits
just tutorials

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.

1. Includes a functional interface and triton backend for 1D/2D/3D. ↩ ↩²