Metadata-Version: 2.4
Name: tinysolve
Version: 0.1.0
Summary: A torch.linalg.solve for tinygrad: batched, differentiable, pivoted linear solve as a single custom kernel
Keywords: tinygrad,linear-algebra,linear-solver,solve,gpu,autograd
Author: FUJISHIGE TEMMA
Author-email: FUJISHIGE TEMMA <tenma.x0@gmail.com>
License-Expression: MIT
License-File: LICENSE
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: Intended Audience :: Developers
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Topic :: Scientific/Engineering
Classifier: Typing :: Typed
Requires-Dist: tinygrad>=0.13.0
Requires-Python: >=3.13
Project-URL: Homepage, https://github.com/FujishigeTemma/tinysolve
Project-URL: Repository, https://github.com/FujishigeTemma/tinysolve
Project-URL: Issues, https://github.com/FujishigeTemma/tinysolve/issues
Description-Content-Type: text/markdown

# tinysolve

A `torch.linalg.solve` for [tinygrad](https://github.com/tinygrad/tinygrad): solve `A @ x = B` for `x`, batched, differentiable, and JIT-friendly — as a **single custom kernel**.

```python
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

```sh
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

```sh
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
