This repository provides a PyTorch implementation of Preconditioned Stochastic Gradient Descent (PSGD), a general-purpose second-order optimizer. It implements various preconditioner structures including Kronecker product, low-rank approximation, dense matrix, affine group, and X-matrix preconditioners. The code offers both Newton-type and whitening preconditioners with multiple update rules for the preconditioner matrices. Key features include support for different geometries (Q0.5EQ1.5, EQ, QEP, etc.), momentum handling, gradient clipping, and numerical stability mechanisms. The library includes optimizers suitable for various deep learning applications as demonstrated in the provided examples (vision transformers, GPT2, LSTMs, CNNs, etc.).
LinearOperators.jl is a Julia library that provides a framework for representing and manipulating linear operators—objects that behave like matrices but are defined by their effect on vectors. It enables lazy composition, transposition, conjugation, and arithmetic (addition, multiplication, scaling) of operators, deferring costly computations until application to a vector. The library includes a rich set of built-in operators: identity, zero, ones, diagonal, restriction/extension, Householder, Hermitian, inverse (via Cholesky, LDL, or general solve), and limited-memory quasi-Newton approximations (L-BFGS and L-SR1) commonly used in optimization. It also provides utility functions for norm estimation, Hermiticity checks, and solving shifted linear systems. The architecture is modular, with a core abstract type hierarchy, separate files for operations, adjoint/transpose wrappers, concatenation, Kronecker product, and timed operators. Extensions support GPU backends (CUDA, AMDGPU, Metal) and automatic differentiation (ChainRulesCore). The code is well-tested, with comprehensive unit tests covering correctness, allocation-free paths, and edge cases. It is primarily used within the JuliaSmoothOptimizers ecosystem for smooth optimization.