cholqr 0.1.0__tar.gz

cholqr-0.1.0/.gitignore

@@ -0,0 +1,5 @@
+__pycache__/
+*.egg-info/
+.venv/
+dist/
+.pytest_cache/

cholqr-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Soki Fukuda
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

cholqr-0.1.0/PKG-INFO

@@ -0,0 +1,111 @@
+Metadata-Version: 2.5
+Name: cholqr
+Version: 0.1.0
+Summary: Fast, compile-friendly shifted CholeskyQR3 for batched tall-skinny matrices in PyTorch
+Project-URL: Homepage, https://github.com/createcentury/cholqr
+Author: Soki Fukuda
+License-Expression: MIT
+License-File: LICENSE
+Keywords: cholesky,gpu,muon,pytorch,qr,randomized-svd
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Mathematics
+Requires-Python: >=3.9
+Requires-Dist: torch>=2.1
+Description-Content-Type: text/markdown
+
+# cholqr
+
+Fast, compile-friendly **shifted CholeskyQR3** for batched tall-skinny matrices in PyTorch.
+It is a drop-in for `torch.linalg.qr(Y, mode="reduced")[0]` when you only need an orthonormal basis of `range(Y)`,
+for example in randomized SVD and sketching, low-rank optimizers such as randomized Muon, or subspace iteration.
+
+```python
+import torch
+from cholqr import cholesky_qr
+
+Y = torch.randn(40, 768, 210, device="cuda")       # batch of tall-skinny matrices
+Q = cholesky_qr(Y)                                 # orthonormal columns, same span as Y
+fast = torch.compile(cholesky_qr, fullgraph=True)  # single graph, no host synchronization
+```
+
+## Why
+
+On GPUs, batched Householder QR (cuSOLVER `geqr2`) processes the batch one matrix at a time.
+
+On an NVIDIA A100, for the sketches of randomized low-rank Muon on nanoGPT (two weight banks, sketch width 210):
+
+| | Householder QR | `cholqr` | full Newton–Schulz (reference) |
+|---|---|---|---|
+| optimizer time per step, both banks, one GPU | 117.6 ms | **15.2 ms** | 17.1 ms |
+
+- `torch.linalg.qr` reaches about **1.6% of FP32 peak** here.
+- On the QR step alone, `cholqr` is 3–14x faster.
+- End to end (1490 steps, A100-80GB), replacing the QR removes about 100 ms per step. Randomized Muon then trains as fast as full Muon at unchanged perplexity.
+
+Library benchmark on an **NVIDIA Tesla T4** (torch 2.14.1, CUDA 13.0; mean of 20 calls; `benchmarks/bench_qr.py`):
+
+| batch x m x r | `torch.linalg.qr` | `cholqr` (compiled) | speedup | `robust=True` | max ‖QᵀQ − I‖_F |
+|---|---|---|---|---|---|
+| 40 x 768 x 210 | 51.7 ms | 10.8 ms | **4.8x** | 45.1 ms | 3.5e-6 |
+| 10 x 3072 x 210 | 23.0 ms | 9.7 ms | **2.4x** | 40.4 ms | 3.5e-6 |
+| 16 x 4096 x 256 | 62.4 ms | 25.5 ms | **2.4x** | 104.2 ms | 4.3e-6 |
+| 4 x 768 x 64 | 1.55 ms | 1.68 ms | 0.9x | 6.2 ms | 1.5e-6 |
+
+- The gain comes from batches of moderately wide sketches. For small problems (e.g. 4 x 768 x 64), Householder QR is already fast, so do not expect a speedup there.
+- `robust=True` is slow on the T4 because the T4 has low FP64 throughput. Use it only for ill-conditioned inputs, preferably on A100/H100-class GPUs.
+- The A100 numbers in the table above come from the randomized-Muon training code (same algorithm).
+
+Run `python benchmarks/bench_qr.py --compile` on your GPU to measure your own shapes.
+
+## How
+
+Shifted CholeskyQR3 (Fukaya et al., 2020) with three implementation choices:
+
+1. **Column scaling before every Gram product.** The shift is scale invariant, and unnormalized inputs do not overflow FP32.
+2. **Explicit r x r triangular inverse applied as one batched GEMM.** This replaces a slow triangular solve over m rows.
+3. **`cholesky_ex` with a per-matrix `torch.where` fallback.** There is no host synchronization, and the routine compiles with zero graph breaks.
+
+## Accuracy and range of validity
+
+| input condition number | default (FP32) | `robust=True` (FP64 first pass) |
+|---|---|---|
+| up to 1e6 | ‖QᵀQ − I‖_F ≈ 1e-5 | ‖QᵀQ − I‖_F ≈ 1e-5 |
+| 5e6 – 1e8 | **breaks down** (orthogonality lost) | ‖QᵀQ − I‖_F ≈ 1e-5, subspace error ≈ 2e-6 |
+
+These numbers were tested on 3072 x 210 inputs with geometric and power-law spectra (`tests/`).
+
+- The default first-pass shift `1e-6 * tr(G)` is a heuristic. The a-priori shift of Fukaya et al. is impractically large in FP32 for large m·r.
+- `robust=True` costs one FP64 Gram product. That is cheap on data-center GPUs (A100/H100) and slow on consumer GPUs.
+- **Real optimizer states** (27 nanoGPT momentum snapshots, sketch width 210, bf16 power iteration as in training):
+  - κ(Y) ≤ 4.5×10⁴.
+  - No fallback was triggered.
+  - The final low-rank update differs from an FP64 Householder reference by at most 7.6×10⁻⁵ (relative).
+  - Sketches computed in bf16 are typically well inside the default range, because bf16 rounding caps their effective condition number.
+- Exactly rank-deficient or zero inputs return finite outputs (not an orthonormal basis of a larger space).
+- Real dtypes only; complex inputs raise `TypeError`. Autograd works (gradients agree with Householder QR to ~1e-7 in tests). Runs on CUDA, CPU and Apple MPS.
+- Status: **experimental (0.x)**. The API may change. Issues and benchmarks on other GPUs are welcome.
+
+```python
+Q = cholesky_qr(Y)               # fast default, cond(Y) <~ 1e6
+Q = cholesky_qr(Y, robust=True)  # ill-conditioned inputs, cond(Y) up to >= 1e8
+```
+
+## Install
+
+```bash
+pip install cholqr        # or: pip install git+https://github.com/createcentury/cholqr
+```
+
+## Citation
+
+Please cite the algorithm:
+
+> T. Fukaya, R. Kannan, Y. Nakatsukasa, Y. Yamamoto, Y. Yanagisawa. *Shifted Cholesky QR for computing the QR
+> factorization of ill-conditioned matrices.* SIAM J. Sci. Comput. 42(1):A477–A503, 2020.
+
+> Y. Yamamoto, Y. Nakatsukasa, Y. Yanagisawa, T. Fukaya. *Roundoff error analysis of the CholeskyQR2 algorithm.*
+> Electron. Trans. Numer. Anal. 44:306–326, 2015.
+
+## License
+
+MIT

cholqr-0.1.0/README.md

@@ -0,0 +1,96 @@
+# cholqr
+
+Fast, compile-friendly **shifted CholeskyQR3** for batched tall-skinny matrices in PyTorch.
+It is a drop-in for `torch.linalg.qr(Y, mode="reduced")[0]` when you only need an orthonormal basis of `range(Y)`,
+for example in randomized SVD and sketching, low-rank optimizers such as randomized Muon, or subspace iteration.
+
+```python
+import torch
+from cholqr import cholesky_qr
+
+Y = torch.randn(40, 768, 210, device="cuda")       # batch of tall-skinny matrices
+Q = cholesky_qr(Y)                                 # orthonormal columns, same span as Y
+fast = torch.compile(cholesky_qr, fullgraph=True)  # single graph, no host synchronization
+```
+
+## Why
+
+On GPUs, batched Householder QR (cuSOLVER `geqr2`) processes the batch one matrix at a time.
+
+On an NVIDIA A100, for the sketches of randomized low-rank Muon on nanoGPT (two weight banks, sketch width 210):
+
+| | Householder QR | `cholqr` | full Newton–Schulz (reference) |
+|---|---|---|---|
+| optimizer time per step, both banks, one GPU | 117.6 ms | **15.2 ms** | 17.1 ms |
+
+- `torch.linalg.qr` reaches about **1.6% of FP32 peak** here.
+- On the QR step alone, `cholqr` is 3–14x faster.
+- End to end (1490 steps, A100-80GB), replacing the QR removes about 100 ms per step. Randomized Muon then trains as fast as full Muon at unchanged perplexity.
+
+Library benchmark on an **NVIDIA Tesla T4** (torch 2.14.1, CUDA 13.0; mean of 20 calls; `benchmarks/bench_qr.py`):
+
+| batch x m x r | `torch.linalg.qr` | `cholqr` (compiled) | speedup | `robust=True` | max ‖QᵀQ − I‖_F |
+|---|---|---|---|---|---|
+| 40 x 768 x 210 | 51.7 ms | 10.8 ms | **4.8x** | 45.1 ms | 3.5e-6 |
+| 10 x 3072 x 210 | 23.0 ms | 9.7 ms | **2.4x** | 40.4 ms | 3.5e-6 |
+| 16 x 4096 x 256 | 62.4 ms | 25.5 ms | **2.4x** | 104.2 ms | 4.3e-6 |
+| 4 x 768 x 64 | 1.55 ms | 1.68 ms | 0.9x | 6.2 ms | 1.5e-6 |
+
+- The gain comes from batches of moderately wide sketches. For small problems (e.g. 4 x 768 x 64), Householder QR is already fast, so do not expect a speedup there.
+- `robust=True` is slow on the T4 because the T4 has low FP64 throughput. Use it only for ill-conditioned inputs, preferably on A100/H100-class GPUs.
+- The A100 numbers in the table above come from the randomized-Muon training code (same algorithm).
+
+Run `python benchmarks/bench_qr.py --compile` on your GPU to measure your own shapes.
+
+## How
+
+Shifted CholeskyQR3 (Fukaya et al., 2020) with three implementation choices:
+
+1. **Column scaling before every Gram product.** The shift is scale invariant, and unnormalized inputs do not overflow FP32.
+2. **Explicit r x r triangular inverse applied as one batched GEMM.** This replaces a slow triangular solve over m rows.
+3. **`cholesky_ex` with a per-matrix `torch.where` fallback.** There is no host synchronization, and the routine compiles with zero graph breaks.
+
+## Accuracy and range of validity
+
+| input condition number | default (FP32) | `robust=True` (FP64 first pass) |
+|---|---|---|
+| up to 1e6 | ‖QᵀQ − I‖_F ≈ 1e-5 | ‖QᵀQ − I‖_F ≈ 1e-5 |
+| 5e6 – 1e8 | **breaks down** (orthogonality lost) | ‖QᵀQ − I‖_F ≈ 1e-5, subspace error ≈ 2e-6 |
+
+These numbers were tested on 3072 x 210 inputs with geometric and power-law spectra (`tests/`).
+
+- The default first-pass shift `1e-6 * tr(G)` is a heuristic. The a-priori shift of Fukaya et al. is impractically large in FP32 for large m·r.
+- `robust=True` costs one FP64 Gram product. That is cheap on data-center GPUs (A100/H100) and slow on consumer GPUs.
+- **Real optimizer states** (27 nanoGPT momentum snapshots, sketch width 210, bf16 power iteration as in training):
+  - κ(Y) ≤ 4.5×10⁴.
+  - No fallback was triggered.
+  - The final low-rank update differs from an FP64 Householder reference by at most 7.6×10⁻⁵ (relative).
+  - Sketches computed in bf16 are typically well inside the default range, because bf16 rounding caps their effective condition number.
+- Exactly rank-deficient or zero inputs return finite outputs (not an orthonormal basis of a larger space).
+- Real dtypes only; complex inputs raise `TypeError`. Autograd works (gradients agree with Householder QR to ~1e-7 in tests). Runs on CUDA, CPU and Apple MPS.
+- Status: **experimental (0.x)**. The API may change. Issues and benchmarks on other GPUs are welcome.
+
+```python
+Q = cholesky_qr(Y)               # fast default, cond(Y) <~ 1e6
+Q = cholesky_qr(Y, robust=True)  # ill-conditioned inputs, cond(Y) up to >= 1e8
+```
+
+## Install
+
+```bash
+pip install cholqr        # or: pip install git+https://github.com/createcentury/cholqr
+```
+
+## Citation
+
+Please cite the algorithm:
+
+> T. Fukaya, R. Kannan, Y. Nakatsukasa, Y. Yamamoto, Y. Yanagisawa. *Shifted Cholesky QR for computing the QR
+> factorization of ill-conditioned matrices.* SIAM J. Sci. Comput. 42(1):A477–A503, 2020.
+
+> Y. Yamamoto, Y. Nakatsukasa, Y. Yanagisawa, T. Fukaya. *Roundoff error analysis of the CholeskyQR2 algorithm.*
+> Electron. Trans. Numer. Anal. 44:306–326, 2015.
+
+## License
+
+MIT

cholqr-0.1.0/benchmarks/bench_qr.py

@@ -0,0 +1,20 @@
+"""Benchmark cholqr.cholesky_qr against torch.linalg.qr on batched tall-skinny matrices (run on a GPU).
+Usage: python benchmarks/bench_qr.py [--shapes 40x768x210,10x3072x210] [--compile]"""
+import argparse, time, torch
+from cholqr import cholesky_qr
+
+ap = argparse.ArgumentParser(); ap.add_argument("--shapes", default="40x768x210,10x3072x210"); ap.add_argument("--compile", action="store_true")
+a = ap.parse_args(); dev = "cuda" if torch.cuda.is_available() else "cpu"
+fn = torch.compile(cholesky_qr, fullgraph=True) if a.compile else cholesky_qr
+def timeit(f, Y, n=20):
+    for _ in range(3): f(Y)
+    if dev == "cuda": torch.cuda.synchronize()
+    t = time.perf_counter()
+    for _ in range(n): f(Y)
+    if dev == "cuda": torch.cuda.synchronize()
+    return (time.perf_counter() - t) / n * 1e3
+for sh in a.shapes.split(","):
+    b, m, r = map(int, sh.split("x")); Y = torch.randn(b, m, r, device=dev)
+    th = timeit(lambda X: torch.linalg.qr(X, mode="reduced")[0], Y); tc = timeit(fn, Y)
+    Q = fn(Y); err = (Q.mT @ Q - torch.eye(r, device=dev)).norm(dim=(-2, -1)).max().item()
+    print(f"{dev} {b}x{m}x{r}: torch.linalg.qr {th:7.2f} ms | cholqr {tc:7.2f} ms | speedup {th/tc:5.1f}x | max ||Q^TQ-I||_F {err:.1e}")

cholqr-0.1.0/benchmarks/results/t4_2026-11-07.txt

@@ -0,0 +1,5 @@
+Tesla T4, torch 2.14.1+cu130 (Colab), mean of 20 calls after 3 warmups
+40x768x210:  torch.linalg.qr 51.66 ms | cholqr eager 11.26 | compiled 10.77 | robust 45.10 ms | speedup 4.8x | max||QtQ-I|| 3.5e-06
+10x3072x210: torch.linalg.qr 22.96 ms | cholqr eager  8.99 | compiled  9.70 | robust 40.44 ms | speedup 2.4x | max||QtQ-I|| 3.5e-06
+4x768x64:    torch.linalg.qr  1.55 ms | cholqr eager  2.61 | compiled  1.68 | robust  6.15 ms | speedup 0.9x | max||QtQ-I|| 1.5e-06
+16x4096x256: torch.linalg.qr 62.35 ms | cholqr eager 23.14 | compiled 25.48 | robust 104.20 ms | speedup 2.4x | max||QtQ-I|| 4.3e-06

cholqr-0.1.0/pyproject.toml

@@ -0,0 +1,21 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "cholqr"
+version = "0.1.0"
+description = "Fast, compile-friendly shifted CholeskyQR3 for batched tall-skinny matrices in PyTorch"
+readme = "README.md"
+license = "MIT"
+requires-python = ">=3.9"
+authors = [{ name = "Soki Fukuda" }]
+dependencies = ["torch>=2.1"]
+keywords = ["qr", "cholesky", "randomized-svd", "muon", "pytorch", "gpu"]
+classifiers = ["Programming Language :: Python :: 3", "Topic :: Scientific/Engineering :: Mathematics"]
+
+[project.urls]
+Homepage = "https://github.com/createcentury/cholqr"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/cholqr"]

cholqr-0.1.0/src/cholqr/__init__.py

@@ -0,0 +1,12 @@
+"""cholqr: fast, compile-friendly shifted CholeskyQR3 for batched tall-skinny matrices in PyTorch.
+
+Replaces ``torch.linalg.qr(Y, mode="reduced")[0]`` when only an orthonormal basis of ``range(Y)`` is needed
+(randomized SVD / sketching, low-rank optimizers such as randomized Muon, subspace iteration).
+
+Algorithm: shifted CholeskyQR3 (Fukaya, Kannan, Nakatsukasa, Yamamoto, Yanagisawa, SIAM J. Sci. Comput. 42(1), 2020)
+built on CholeskyQR2 (Yamamoto, Nakatsukasa, Yanagisawa, Fukaya, ETNA 44, 2015).
+"""
+from .core import cholesky_qr, orthonormal_basis
+
+__all__ = ["cholesky_qr", "orthonormal_basis"]
+__version__ = "0.1.0"

cholqr-0.1.0/src/cholqr/core.py

@@ -0,0 +1,68 @@
+import torch
+
+__all__ = ["cholesky_qr", "orthonormal_basis"]
+
+
+def cholesky_qr(Y: torch.Tensor, shift: float = 1e-6, passes: int = 3, out_dtype=None,
+                robust: bool = False) -> torch.Tensor:
+    """Orthonormal basis Q of range(Y) for a (batch of) tall-skinny matrix Y of shape (..., m, r), r <= m.
+
+    Shifted CholeskyQR3 with three implementation choices that make it fast on GPUs:
+
+    * every pass rescales the columns before forming the Gram matrix (scale invariant, no FP32 overflow);
+    * the r x r triangular inverse is formed explicitly and applied as one batched GEMM (a triangular solve
+      over m rows is slow on GPUs);
+    * ``torch.linalg.cholesky_ex`` is used and a failed pass is skipped per matrix with ``torch.where``, so the
+      routine never synchronizes with the host and compiles into a single graph under ``torch.compile``.
+
+    Pass 1 uses the shift ``shift * tr(G)`` (a heuristic, not the a-priori shift of Fukaya et al., whose FP32 value is
+    impractically large for large m*r); later passes use ``(u / r) * tr(G)`` with u the FP32 unit roundoff, restoring
+    ``||Q^T Q - I||`` to O(u). Computation is in FP32; the result is cast to ``out_dtype`` (default: the dtype of Y).
+
+    Range of validity (tested on 3072 x 210 inputs with geometric and power-law spectra):
+
+    * default (FP32): reliable for cond(Y) up to about 1e6. Beyond ~5e6 the factorization breaks down
+      (orthogonality is lost and passes are skipped). Inputs that come from bf16 sketches are usually fine further,
+      because bf16 rounding caps their effective condition number.
+    * ``robust=True``: the first pass runs in FP64 with a tiny shift, which handles cond(Y) up to at least 1e8 with
+      ``||Q^T Q - I||_F ~ 1e-5``. It costs one FP64 Gram product (cheap on data-center GPUs, slow on consumer GPUs).
+
+    Args:
+        Y: tensor of shape (..., m, r) with r <= m.
+        shift: relative shift of the first pass.
+        passes: number of CholeskyQR passes (3 = shifted CholeskyQR3).
+        out_dtype: dtype of the result.
+        robust: run the first pass in FP64 (see above).
+    Returns:
+        Q of shape (..., m, r) with orthonormal columns spanning range(Y) (when Y has full column rank).
+    """
+    if Y.dim() < 2:
+        raise ValueError("Y must have at least 2 dimensions (..., m, r)")
+    if Y.is_complex():
+        raise TypeError("cholesky_qr supports real tensors only (got a complex dtype)")
+    m, r = Y.shape[-2], Y.shape[-1]
+    if r > m:
+        raise ValueError(f"cholesky_qr expects tall-skinny input (r <= m), got m={m}, r={r}")
+    u = torch.finfo(torch.float32).eps
+    Q = Y.double() if robust else Y.float()
+    for i in range(passes):
+        if robust and i == 1:
+            Q = Q.float()
+        eye = torch.eye(r, device=Q.device, dtype=Q.dtype)
+        Q = Q / Q.norm(dim=-2, keepdim=True).clamp_min(1e-30)
+        G = Q.mT @ Q
+        tr = G.diagonal(dim1=-2, dim2=-1).sum(dim=-1)[..., None, None]
+        if i == 0:
+            s = (1e-14 if robust else shift) * tr
+        else:
+            s = (u / r) * tr
+        L, info = torch.linalg.cholesky_ex(G + (s + 1e-30) * eye)
+        L_inv = torch.linalg.solve_triangular(L, eye.expand_as(L), upper=False)
+        Q_next = Q @ L_inv.mT
+        Q = torch.where((info == 0)[..., None, None], Q_next, Q)
+    return Q.float().to(Y.dtype if out_dtype is None else out_dtype)
+
+
+def orthonormal_basis(Y: torch.Tensor, **kwargs) -> torch.Tensor:
+    """Alias of :func:`cholesky_qr` (drop-in for ``torch.linalg.qr(Y, mode="reduced")[0]`` when only Q is used)."""
+    return cholesky_qr(Y, **kwargs)

cholqr-0.1.0/tests/test_cholqr.py

@@ -0,0 +1,95 @@
+import math
+import pytest
+import torch
+from cholqr import cholesky_qr
+
+
+def make_spectrum(m, r, kappa, kind, seed=0):
+    g = torch.Generator().manual_seed(seed)
+    U = torch.linalg.qr(torch.randn(m, r, generator=g, dtype=torch.float64))[0]
+    V = torch.linalg.qr(torch.randn(r, r, generator=g, dtype=torch.float64))[0]
+    if kind == "geometric":
+        s = torch.logspace(0, -math.log10(kappa), r, dtype=torch.float64)
+    else:
+        s = torch.arange(1, r + 1, dtype=torch.float64) ** -(math.log(kappa) / math.log(r))
+    return ((U * s) @ V.mT).float()
+
+
+def orth_err(Q):
+    Q = Q.double()
+    return (Q.mT @ Q - torch.eye(Q.shape[-1], dtype=torch.float64)).norm(dim=(-2, -1)).max().item()
+
+
+def subspace_err(Q, Y):
+    Qr = torch.linalg.qr(Y.double())[0]
+    return torch.linalg.matrix_norm(Q.double() @ Q.double().mT - Qr @ Qr.mT, ord=2).max().item()
+
+
+@pytest.mark.parametrize("kind", ["geometric", "powerlaw"])
+@pytest.mark.parametrize("kappa", [1e1, 1e3, 1e5, 1e6])
+def test_default_fp32_up_to_1e6(kind, kappa):
+    Y = make_spectrum(3072, 210, kappa, kind)
+    Q = cholesky_qr(Y)
+    assert orth_err(Q) < 5e-5
+    assert subspace_err(Q, Y) < 1e-2 * max(1.0, kappa / 1e5)
+
+
+@pytest.mark.parametrize("kind", ["geometric", "powerlaw"])
+@pytest.mark.parametrize("kappa", [1e6, 1e7, 1e8])
+def test_robust_up_to_1e8(kind, kappa):
+    Y = make_spectrum(3072, 210, kappa, kind)
+    Q = cholesky_qr(Y, robust=True)
+    assert orth_err(Q) < 5e-5
+    assert subspace_err(Q, Y) < 1e-4
+
+
+def test_default_breaks_down_beyond_range():
+    # documents the limitation: plain FP32 loses orthogonality for cond ~ 1e7
+    Y = make_spectrum(3072, 210, 1e7, "geometric")
+    assert orth_err(cholesky_qr(Y)) > 1.0
+
+
+def test_batched_matches_unbatched():
+    Y = torch.stack([make_spectrum(768, 64, 1e3, "powerlaw", seed=k) for k in range(4)])
+    Qb = cholesky_qr(Y)
+    for k in range(4):
+        assert torch.allclose(Qb[k], cholesky_qr(Y[k]), atol=1e-5)
+
+
+def test_degenerate_inputs_are_finite():
+    Z = torch.zeros(64, 8)
+    assert torch.isfinite(cholesky_qr(Z)).all()
+    Y = torch.randn(64, 4) @ torch.randn(4, 8)                       # exactly rank 4 < 8
+    assert torch.isfinite(cholesky_qr(Y)).all()
+
+
+def test_dtypes():
+    Y = torch.randn(128, 16)
+    assert cholesky_qr(Y.bfloat16()).dtype == torch.bfloat16
+    assert cholesky_qr(Y, out_dtype=torch.float64).dtype == torch.float64
+
+
+def test_rejects_wide_input():
+    with pytest.raises(ValueError):
+        cholesky_qr(torch.randn(8, 16))
+
+
+@pytest.mark.parametrize("robust", [False, True])
+def test_compiles_without_graph_breaks(robust):
+    explain = torch._dynamo.explain(cholesky_qr)(torch.randn(256, 32), robust=robust)
+    assert explain.graph_break_count == 0
+
+
+def test_rejects_complex():
+    with pytest.raises(TypeError):
+        cholesky_qr(torch.randn(64, 8, dtype=torch.complex64))
+
+
+def test_autograd_matches_householder():
+    torch.manual_seed(0)
+    Y = torch.randn(64, 8, dtype=torch.float64, requires_grad=True)
+    W = torch.randn(64, 64, dtype=torch.float64)
+    loss = lambda Q: ((Q @ Q.mT) * W).sum()      # sign-invariant
+    g1, = torch.autograd.grad(loss(cholesky_qr(Y, out_dtype=torch.float64)), Y)
+    g2, = torch.autograd.grad(loss(torch.linalg.qr(Y)[0]), Y)
+    assert ((g1 - g2).norm() / g2.norm()).item() < 1e-5