tinysolve 0.1.0

A torch.linalg.solve for tinygrad: batched, differentiable, pivoted linear solve as a single custom kernel

tinysolve

A torch.linalg.solve for tinygrad: solve A @ x = B for x, batched, differentiable, and JIT-friendly — as a single custom kernel.

from tinygrad import Tensor
from tinysolve import solve

A = Tensor.randn(2000, 16, 16, dtype="float64") + 16 * Tensor.eye(16, dtype="float64")
B = Tensor.randn(2000, 16, 4, dtype="float64")

x = solve(A, B)  # (2000, 16, 4)

# differentiable (higher-order too)
gA, gB = solve(A, B).sum().gradient(A, B)

Install

uv add tinysolve  # or: pip install tinysolve

API

solve(A, B) follows torch.linalg.solve shape semantics:

• A: (*, n, n) — a (batch of) square matrices; batch dims broadcast against B's.
• B: (*, n, k) for a matrix RHS → result (*, n, k), or a batch of vectors (B.ndim == 1, or B.ndim == A.ndim - 1 with B.shape == A.shape[:-1]) → result (*, n).

Partial pivoting is used, so matrices like [[0, 1], [1, 0]] are handled robustly (unpivoted LU would produce nan). Gradients use the implicit-function identity for a linear solve (same as PyTorch), and the backward reuses the same differentiable solve, so higher-order gradients work.

Contract difference vs torch: this kernel is lazy/JIT-compatible, so it cannot raise on a singular A without forcing a device sync. A singular input surfaces as non-finite (inf/nan) output instead of a LinAlgError.

Why a custom kernel

Batched small-matrix solves are memory-bound (arithmetic intensity ≈ n/12 FLOP/byte, far below the fp64 roofline ridge) — the game is how few times each matrix streams through HBM. A high-level tensor program can't win here: every op round-trips global memory, and the sequential elimination dependency crosses kernel boundaries (~64 kernels per solve for a tensor-level QR).

tinysolve instead uses tinygrad's custom_kernel (UOp) API with the same execution model cuSOLVER uses for small batches: one thread per matrix, the whole n×n system held in registers, and a fully-unrolled partial-pivoted Gaussian elimination + back-substitution run entirely in-register — one kernel, ~one HBM pass per matrix.

Performance (H100, fp64, JIT, vs torch.linalg.solve, k=1)

speedup vs torch	M=500	M=2000	M=20000
n=4	4.40×	4.46×	4.62×
n=8	3.64×	3.60×	2.81×
n=16	0.74×	0.73×	0.89×

3–4× faster than torch for n ≤ 8, ~parity at n = 16. Beyond n ≈ 16 the per-thread working set (n·(n+k) doubles) exceeds the register file and spills; cuSOLVER's sub-warp-per-matrix cooperative kernels win there.

Development

uv run pytest                     # tests (property-based, vs torch as oracle)
uv run pytest tests/benchmark.py  # benchmarks
uv run ruff check . && uv run ty check

License

MIT