torch_vn_algebra: block-diagonal operators, symmetry sectors and quantum channels in PyTorch

torch_vn_algebra is a Python library for batched computation with block-diagonal matrices in symmetry sectors, the finite-dimensional Type I von Neumann algebras ⊕ M_k(ℂ), in PyTorch on CPU and GPU. It provides random operators and Haar-random unitaries, density matrices, quantum channels (Kraus, Choi, Stinespring), Schrödinger and Lindblad dynamics, and spin chains in U(1), momentum and SU(2) sectors; it computes the Wedderburn (block-diagonal) decomposition and the commutant of the algebra generated by given matrices or by a group representation, conditional expectations and Jones indices, Jordan algebras and quantum Rényi divergences. Operators, states and channels are batched over thousands of samples, differentiable with autograd, and available in single or double precision.

It is written for physicists (superselection and charge sectors, exact diagonalisation, open quantum systems), quantum-information researchers (states, channels, entropies, error correction), mathematicians (operator algebras, random matrices, representation theory) and machine-learning researchers (equivariant linear layers as commutants). The theory behind every function is in self-contained lecture notes with complete proofs, and the library is shown at work in a tutorial whose code is run by the tests, both in English, Spanish, French, German, Chinese, Japanese and Russian; documentation website: https://andre-a-no.github.io/von_neumann_type_i/.

The algebras are direct sums of full matrix algebras:

\mathcal M = \bigoplus_{c=1}^{C} M_{k_c}(\mathbb C)\quad\text{acting on}\quad \mathcal H=\bigoplus_{c=1}^{C}\mathbb C^{k_c}.

(Nominal factor sizes n_factors $n_c\ge k_c$ describe computations in the corner $p\mathcal M_0p$ of a larger algebra $\mathcal M_0=\bigoplus_c M_{n_c}$ with rank-$k_c$ projections $p_c$; everything refers to the corner.)

Operators are stored as one tensor of shape (batch, C, k_max, k_max): the batch axis holds samples (or optimisation starts), the sector axis holds the direct summands, and block c occupies the top-left k_c × k_c corner (the rest is zero padding).

Companion paper: I. Nikolaeva, A. Novikov, Finite-Dimensional Type I von Neumann Algebras in PyTorch: A GPU-Accelerated Framework for Random Block-Diagonal Operators, arXiv:2606.15882 (source and PDF in man/article/).

This README is the entry point of the repository:

Features

• Block-diagonal operators in batches – an operator of ⊕c M{k_c} is one tensor (batch, C, k_max, k_max); products, adjoints, traces, norms and functional calculus run as batched kernels on CPU or CUDA GPU (Tutorial 1, 2.5).
• Random matrices and Haar-random unitaries – Haar measure on U(n), O(n), SU(n), SO(n), COE, CSE, random operators with a prescribed spectrum, GUE/GOE level statistics, Weingarten-calculus checks (Tutorial 5).
• Quantum states and entropy – density matrices (Hilbert–Schmidt, Bures, fixed rank), Gibbs states, von Neumann entropy, relative entropy, fidelity, trace distance (Tutorial 3).
• Quantum channels – Kraus operators, Choi matrix, Stinespring dilation, superoperator; dephasing, depolarizing, amplitude damping, random channels, channels between sectors (Tutorial 6).
• Tensor products and entanglement – partial trace, fusion of U(1) charges, entanglement entropy, Page's formula (Tutorial 7).
• Schrödinger and Lindblad (GKSL) master equations – exact propagators, RK4, time steps that are CPTP channels, quantum jumps between sectors, batched Hamiltonians (Tutorial 8).
• Spin chains and exact diagonalisation in symmetry sectors – XXZ and Heisenberg chains in U(1) magnetisation, momentum and SU(2) sectors, Lanczos ground states, Krylov time evolution, disorder averages, level-spacing ratio (Tutorial 9).
• Wedderburn decomposition, commutants and group representations – block diagonalisation of the matrix *-algebra generated by matrices, a Lie algebra or a finite group (isotypic decomposition, multiplicities, Schur–Weyl duality), real/complex/quaternionic type, Clebsch–Gordan coefficients and the Wigner–Eckart theorem (Tutorial 4, Appendix B).
• Quantum error correction – noiseless subsystems, decoherence-free subspaces, the Knill–Laflamme conditions, operator-algebra error correction (Tutorial 10).
• Conditional expectations and inclusions of algebras – Jones index, basic construction, Pimsner–Popa constants, Bratteli diagrams (Tutorial 2.6, example_inclusions.py).
• Quantum Rényi divergences – Petz and sandwiched Rényi divergences, Umegaki relative entropy, D_max, Petz recovery maps (example_divergences.py).
• Constrained optimisation – batched multistart optimisation over unitaries and positive operators, trace inequalities (Tutorial 11).
• Equivariant machine learning – equivariant linear layers as the commutant of a group representation (Tutorial 12).
• Jordan algebras – Jordan product, quadratic representation, spin factors and the Lorentz cone, the Jordan algebra generated by matrices (Tutorial 13).

When to use it, and alternatives

torch_vn_algebra pays off when the block structure is part of the problem: many sectors of different sizes (superselection or charge sectors, isotypic components of a symmetry), many samples or optimisation starts at once, and quantities that are defined by the algebra (its traces, states, centre, conditional expectations, channels between sectors). For one dense matrix at a time, for very large sparse systems, or for tensor networks, the established tools below are the better choice; the sections What to compare with of the tutorial compare them in detail.

Features in detail

• Random operators with any spectrum – operator_from_eigenvalues(sampler, ...) builds A_c = U_c diag(λ) U_c* from a user callable sampler(dim) -> (dim,) or (batch, dim).
• Unitary ensembles – Haar on U(n)/O(n), SU(n)/SO(n), COE, CSE, random diagonal phases; batched (random_unitary(n, measure, batch_size)), and block-diagonal unitaries in the algebra (random_unitary_operator).
• Lazy evaluation – X @ Y, X + Y, X * s and functional calculus build a recipe that is computed when .matrix (or a scalar functional) is requested; methods that need a property of the base operator first (e.g. X.inv, X.sqrt() checking positivity) evaluate the base.
• Functional calculus at four levels (module functional) – apply_function(f) for self-adjoint and for normal operators (unitaries; via the Hermitian eigensolver of H + αK), abs, inverse (pseudo-inverse with a relative cutoff), polar, norms and singular values (on GPUs for blocks > 32 via the Hermitian eigensolver of [[0, A], [A*, 0]], since batched SVD stops at 32), holomorphic_function(f) for any operator (contour integral, batched LU solves, error estimate); sqrt, log, power, expm, entropy; eigenvalues() without eigenvectors. Gradients are finite at degenerate spectra (Daleckii–Krein divided differences). Opt-in cache of the eigendecomposition: A.cache_spectrum() or with torch_vn_algebra.caching(): ....
• Three trace functionals – Tr_blunt (Σ_c Tr A_c), Tr_norm (Σ_c Tr A_c / k_c) and the tracial state tau_vN ((1/C) Σ_c Tr A_c / k_c); norms trace_norm, frobenius_norm, operator_norm; michelson_contrast.
• Hilbert space utilities – HilbertSpace with inner products, orthonormal bases (standard / random / Haar), embedding and restriction of vectors and operators.
• Inclusions of algebras (module inclusions) – Inclusion(N, M, Lambda) for N ⊂ M given by an inclusion matrix (and an optional batched unitary): the trace-preserving conditional expectation for any faithful trace, Markov trace and index ‖Λ‖², Watatani index, Pimsner–Popa constants, basic construction, Jones projection and tower, relative commutant, Bratteli diagrams, and the U(1) spin-chain tower A_L ⊂ A_{L+1} (index 4cos²(π/(2L+4))).
• Wedderburn decomposition (module decompose) – the von Neumann algebra generated by given matrices or by a finite group, V*AV = ⊕c M{k_c} ⊗ 1_{m_c}, by the randomised algorithm of Murota et al. (2010), verified before it is returned; maps between n×n matrices and operators of ⊕c M{k_c} and of the commutant ⊕c M{m_c}.
• SU(2) symmetry of spin-1/2 chains (module su2) – SU2Chain: the algebra ⊕S M{m_S} of SU(2)-invariant operators, orthonormal highest-weight (Yamanouchi) bases in the U(1) sectors of SpinChain, Heisenberg Hamiltonians with J₂ and fields, full spectra with quantum numbers, exact Clebsch–Gordan coefficients, Wigner–Eckart reduced matrix elements and spin correlations (L = 14: largest block 1001 instead of 3432).
• Quantum Rényi divergences (module divergences) – Petz and sandwiched Rényi divergences with exact support conventions, Umegaki, D_max and D_0 as limits, per-sector decomposition, Petz, rotated and universal recovery maps, noncommutative L^p and Kosaki norms; batched and differentiable.
• Jordan algebras (module jordan) – Jordan product, multiplication operator, quadratic representation P(y)x = yxy, Jordan triple product and powers, the transpose as a Jordan automorphism, the Jordan algebra generated by Hermitian matrices (jordan_algebra_generated, compared with the *-algebra they generate), spin factors V_m and the Lorentz cone (SpinFactor).

Installation

Requirements: Python ≥ 3.8, PyTorch ≥ 1.13, NumPy ≥ 1.21 (the whole test suite runs in CI with exactly these minimal versions and with the latest ones on Python 3.8–3.13); SciPy ≥ 1.7 only for the tests and scripts. On Python 3.8 install the wheel from PyPI (pip install torch_vn_algebra); building from source needs Python ≥ 3.9.

From PyPI (install PyTorch for your CPU/GPU first, see https://pytorch.org):

pip install torch_vn_algebra

From the repository (needed for the scripts, tests and the paper):

git clone https://github.com/andre-a-no/von_neumann_type_i.git
cd von_neumann_type_i
pip install -e .              # library only (torch, numpy)
pip install -e ".[scripts]"   # + scipy, pandas, tqdm, matplotlib, seaborn, mpmath for scripts/
pip install -e ".[test]"      # + pytest

Quick start

import torch
from torch_vn_algebra import TypeIAlgebra

device = 'cuda' if torch.cuda.is_available() else 'cpu'

# M = M_4(C) (+) M_6(C), acting on C^4 (+) C^6
alg = TypeIAlgebra(n_factors=[4, 6], k_factors=[4, 6], device=device)

# 100 random positive operators: eigenvalues ~ U[0, 1], Haar-random eigenvectors per block
X = alg.operator_from_eigenvalues(lambda dim: torch.rand(100, dim, device=device),
                                  batch_size=100, force_positive=True, force_self_adjoint=True)

print(X.Tr_blunt().real.mean())            # sum of block traces
print(X.Tr_norm().real.mean())             # sum of normalised block traces
print(X.tau_vN().real.mean())              # tracial state (1/C) sum_c Tr(A_c)/k_c (complex in general)
print(X.michelson_contrast.mean())         # (lmax - lmin)/(lmax + lmin), shape (100,)

# operations stay lazy until .matrix or a scalar functional is requested
U = alg.random_unitary_operator(batch_size=100)        # block-diagonal Haar unitary
Z = X.sqrt() @ U @ X.sqrt()
print(Z.trace_norm().mean(), X.entropy().mean())

Channels, states and dynamics

import torch
from torch_vn_algebra import TypeIAlgebra, channels, dynamics, states

alg = TypeIAlgebra([4, 4], [4, 4], complex_valued=True)          # two superselection sectors
rho = states.random_density_matrix(alg, batch_size=100)            # HS-random states on M
H = alg.operator_from_eigenvalues(lambda d: torch.randn(100, d), batch_size=100,
                                  force_self_adjoint=True)

Phi = channels.random_channel(alg, kraus_rank=2, batch_size=100)   # Haar Stinespring isometry
print(Phi.is_trace_preserving(), Phi.is_unital())
print(states.relative_entropy(Phi(rho), Phi(states.gibbs_state(H, 1.0))))

psi0 = torch.zeros(100, 2, 4, dtype=torch.complex64)               # state vectors, block layout
psi0[:, 0, 0] = 1.0
psi_t = dynamics.schrodinger(psi0, H, times=[0.0, 1.0, 2.0])       # exact, shape (3, 100, 2, 4)

N = alg.from_blocks([torch.diag(torch.arange(4.0))] * 2)           # dephasing jump operator
rho_t = dynamics.lindblad_evolve(rho, H, jumps=[N], times=[0, 1, 2])  # RK4, list of states
Lt = dynamics.lindblad_channel(alg, H, jumps=[N], t=2.0)            # exact exp(tL) as a Channel
print((Lt(rho) - rho_t[-1]).frobenius_norm().max())

Inclusions, symmetries and divergences

import torch
from torch_vn_algebra import TypeIAlgebra, SU2Chain, states
from torch_vn_algebra import inclusions as inc, decompose as dc, divergences as dv

# inclusion N = M_2 (+) C  c  M = M_4 (+) M_3 with inclusion matrix Lambda
N = TypeIAlgebra([2, 1], [2, 1], precision='double')
M = TypeIAlgebra([4, 3], [4, 3], precision='double')
I = inc.Inclusion(N, M, [[1, 2], [1, 1]])
E = I.conditional_expectation('markov')                  # the trace-preserving conditional expectation M -> N
y = M.operator(torch.randn(5, M.C, M.k_max, M.k_max, dtype=M.dtype))
print(E(y).algebra is N)                                 # True: E(y) is an operator of N (batch 5)
print(I.index(), E.pimsner_popa_constant())              # [M : N] = ||Lambda||^2, Pimsner-Popa constant
print([round(T.index(), 6) for T in inc.chain_tower(4)])  # spin-chain tower A_L c A_{L+1}: 4 cos^2(pi/(2L+4))

# the algebra generated by a symmetry: S_3 permuting three qubits
G = dc.permutation_representation(3, 2)                  # (6, 8, 8): all group elements
D = dc.from_group(G)
print(D.k_factors, D.multiplicities)                     # (1, 2), (4, 2): pi(S_3)'' = M_1 (+) M_2
blocks = D.to_algebra(G)                                 # the group elements as operators of D.algebra

# SU(2) blocks of the Heisenberg ring, and Renyi divergences of states on M
chain = SU2Chain(12, 'periodic', complex_valued=False)
print(chain.multiplicities, chain.ground_state(chain.heisenberg()).energy)   # [132, 297, ...], -5.387390917
rho, sigma = states.random_density_matrix(M, 10), states.random_density_matrix(M, 10)
print(dv.sandwiched_renyi_divergence(rho, sigma, 2.0), dv.petz_renyi_divergence(rho, sigma, 0.5))

Lazy operators are evaluated once and cached; after that they drop their recipe, so the intermediate results of an expression such as X @ U @ Y can be garbage-collected. Laziness defers and skips work but does not fuse operations; where it matters the library uses direct formulas instead (e.g. Tr(ρA) in O(k²) via states.trace_of_product).

Channels map each sector to itself, so the Heisenberg dual is the same for Tr_blunt, Tr_norm and tau_vN, and Hamiltonian / Lindblad dynamics with generators in M conserve the sector probabilities Tr rho_c. Unitary dynamics requires complex_valued=True.

Number field and precision are two independent switches of every algebra: complex_valued (real / complex) and precision ('single' / 'double'), i.e. float32, float64, complex64 or complex128; alg.like(complex_valued=..., precision=...) and op.cast(alg2) move between them, alg.bytes_per_operator(B) reports the memory. Rules of thumb (details and measurements in the paper, scripts/numerics.py):

complex_valued=False gives real algebras (orthogonal instead of unitary groups); complex input with a non-zero imaginary part is rejected there. Each force_* flag (self_adjoint, positive, normal, invertible, projection) tags the operator after checking it against the sampled eigenvalues. The batch_size argument of TypeIAlgebra only sets the batch of its Hilbert-space bases; every factory takes its own batch_size.

More in examples/: random Hamiltonian with a parity symmetry, free additive convolution, a trace inequality, unitary ensembles, a Zipf density matrix, decoherence inside superselection sectors, the Michelson contrast under unital and non-unital channels, the spectral form factor of GUE vs. Poisson Hamiltonians, inclusions and their indices (example_inclusions.py), Schur–Weyl duality from the symmetries alone (example_decompose.py), SU(2) multiplets of the Heisenberg ring (example_su2.py), Rényi divergences with data processing and recovery (example_divergences.py), and Jordan algebras and spin factors (example_jordan.py).

Reproducing the paper

The paper (man/article/main.tex, compiled man/article/main.pdf) contains no number typed by hand: every table and figure comes from the scripts in scripts and the listings in man/article/listings, run by one driver:

bash scripts/run_all.sh                    # quick mode, about 4 minutes on a 4-thread CPU
bash scripts/run_all.sh full cuda          # sizes of the paper, on a GPU
cd man/article && pdflatex main.tex && pdflatex main.tex

Running the full experiments on a GPU server

git clone https://github.com/andre-a-no/von_neumann_type_i && cd von_neumann_type_i
git checkout claude/funny-ptolemy-confo7          # until it is merged into main
python -m venv .venv && source .venv/bin/activate
pip install torch                                  # the CUDA build matching the server's driver, see pytorch.org
pip install -e ".[scripts,test]"
python -m pytest -q                                # ~30 s, must pass
python scripts/gpu_check.py                        # ~10 min: correctness against the CPU, no per-sector work, timings
bash scripts/run_all.sh check cuda           # full sizes, minimal repetitions: tests every code path and the GPU memory
nohup bash scripts/run_all.sh full cuda > full_run.log 2>&1 &   # the full run (or inside tmux)

• Scripts and the parts of the paper they produce: validate_known_results (Section 5), numerics (Section 3), inequality_search and bounds_search (Section 6), renyi_threshold (Section 6.1: the data-processing thresholds of the sandwiched and Petz Rényi divergences, located by batched optimisation with violations certified in mpmath; about 2 minutes in quick mode on one CPU thread), chains (Section 7), baselines (Section 9), benchmark_library (Section 10).
• The driver first prints the GPU, its memory, CUDA and cuDNN versions (also saved to man/article/generated/logs/environment.txt) and stops if CUDA is not available.
• Every script writes its own log to man/article/generated/logs/<script>.log, including run-time estimates from the cost module. A failing script does not stop the others; the run ends with a summary.
• ONLY="chains baselines" bash scripts/run_all.sh full cuda reruns a subset.
• scripts/gpu_check.py (about 10 minutes, most of it the CPU side of the timings) has three parts: (1) correctness: the same inputs on the CPU and the GPU, all operations of the paper, single and double precision, Haar sampling (Gram-Schmidt on the GPU against QR, moments); (2) structure: host synchronisations and kernel launches per operation must not grow with the number of sectors (C = 1, 16, 256 at a fixed total size); (3) timings against all CPU cores, with the verdict on the hypothesis "the GPU is faster for large batches" and a comparison of the batched generation with the per-sector algorithm of the first full run. It exits with an error if (1) or (2) fails, and writes gpu_check.json. Run it first on a new machine or after changing the library. tests/test_sync_free.py checks the absence of host reads on the CPU, in the normal test suite.
• Duration. On an RTX 5090 (16-thread CPU) the first full run took about 5 hours: validate 12 min, numerics and baselines 3 min, benchmark 8 min, chains 7 min, inequality_search 3.4 h, bounds_search 1.3 h; most of inequality_search was Haar sampling with a per-matrix QR on the GPU, now batched. On a 4-thread CPU the check run takes about 33 minutes (validate 90 s, numerics 88 s, baselines 45 s, benchmark 189 s, chains 623 s, inequality_search 49 s, bounds_search 871 s; every log ends with a [resources] line).
• Interruptions. inequality_search and bounds_search save every finished part to generated/partial/; after a crash or reboot, rerunning the same command continues where it stopped and gives the same numbers.
• Several GPUs. The scripts are independent, e.g. CUDA_VISIBLE_DEVICES=0 ONLY=bounds_search bash scripts/run_all.sh full cuda and CUDA_VISIBLE_DEVICES=1 ONLY="validate_known_results numerics baselines benchmark_library chains inequality_search" bash scripts/run_all.sh full cuda in two terminals (the examples run only when ONLY is not set).
• ... full cuda tf32 repeats the run with TF32 tensor-core matrix products (Ampere or newer), to measure their effect; the default is true FP32.
• Results land in man/article/generated/ (tables as .tex and .json) and man/article/figures/. Commit those two directories and the logs, then rebuild the PDF; the tables state the device and mode they came from.
• Memory. The largest single step is the dense diagonalisation of a 12870-dimensional sector in chains (about 10 GB on the CPU in the check run); use a GPU with at least 24 GB, preferably 40-80 GB, and a host with at least 16 GB RAM. The check run shows the peak GPU memory of every script in its [resources] line.
• Double precision is used throughout the validation and chain scripts; a data-centre GPU (A100, H100) is much faster there than a consumer GPU (FP64 at 1/2 versus 1/32-1/64 of the FP32 rate).

Repository layout

torch_vn_algebra/    library: algebra.py (TypeIAlgebra, Operator), hilbert_space.py, channels.py,
                     states.py, dynamics.py, composite.py, chains.py, krylov.py, optimize.py, cost.py,
                     functional.py, inclusions.py, decompose.py, su2.py, divergences.py,
                     jordan.py
tests/               pytest suite (CPU, runs in CI)
examples/            short usage examples of the library
man/                 the documentation:
  lecture_notes/     the theory: definitions, theorems, complete proofs (7 languages)
  tutorial/          the library in the context of the theory, executed by the tests (7 languages)
  article/           the paper: main.tex, main.pdf, generated/ (tables), figures/, listings/ (its programs)
scripts/             run_all.sh and the scripts behind every table and figure of the paper

Tests

pytest -q

Contributing

Bug reports, questions and pull requests are welcome: see CONTRIBUTING.md and the Code of Conduct. Changes between versions are listed in CHANGELOG.md.

Limitations

• SVD dominates the cost for k_max ≳ 200 with large batches.
• SpinChain(2, 'periodic') has a single bond (the wrap bond is not counted twice), so it equals the open chain.
• Channel, Hamiltonians and Lindblad generators preserve the sectors; maps between sectors are InterSectorChannels. Only abelian charges are fused. Of the non-abelian symmetries, SU(2) reduces the blocks of spin-1/2 chains (SU2Chain, dense up to L ≈ 16); SU(2) combined with momentum and chains of larger spins are not implemented. decompose handles any symmetry given by matrices, but densely (n up to a few hundred) and over C only. Operators on all sectors of a chain are stored padded (work in single sectors for large chains).
• Sandwiched Rényi divergences for small α < 1 of ill-conditioned states are not accurate (the sandwich σ^g ρ σ^g, g = (1-α)/(2α), cannot be resolved in floating point).
• lindblad_channel, choi and superoperator work with k_c² × k_c² matrices per block, which limits them to blocks of a few tens; use lindblad_evolve for larger blocks.
• Gradients of the functional calculus are finite at degenerate spectra; gradients through the eigenvectors returned by eigh are not. Second derivatives of f(A) are correct at non-degenerate spectra.
• holomorphic_function needs a circle that encloses the spectrum and on which f is holomorphic (the default circle suits entire functions such as exp); for self-adjoint or normal operators apply_function is exact and cheaper.
• On GPUs, batched Hermitian eigensolvers cover all block sizes only in recent PyTorch/cuSOLVER releases.
• Finite dimensions and Type I only.

Citation

Please cite the paper and, for reproducibility, the version of the software you used (metadata in CITATION.cff; GitHub shows a "Cite this repository" button; every release is archived with a DOI on Zenodo).

@misc{nikolaeva_novikov_2026_torch_vn_algebra,
  title  = {Finite-Dimensional Type I von Neumann Algebras in PyTorch:
            A GPU-Accelerated Framework for Random Block-Diagonal Operators},
  author = {Nikolaeva, Irina and Novikov, Andrej},
  year   = {2026},
  eprint = {2606.15882},
  archivePrefix = {arXiv},
  primaryClass  = {cs.MS},
  url    = {https://arxiv.org/abs/2606.15882}
}

@software{torch_vn_algebra,
  title   = {torch\_vn\_algebra: block-diagonal operators, symmetry sectors and quantum channels in {PyTorch}},
  author  = {Nikolaeva, Irina and Novikov, Andrej},
  year    = {2026},
  version = {1.0.0},
  url     = {https://github.com/andre-a-no/von_neumann_type_i}
}

License

MIT, see LICENSE.