torchsparsegradutils 0.2.5__tar.gz → 0.2.6__tar.gz

{torchsparsegradutils-0.2.5/torchsparsegradutils.egg-info → torchsparsegradutils-0.2.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torchsparsegradutils
-Version: 0.2.5
+Version: 0.2.6
 Summary: A collection of utility functions to work with PyTorch sparse tensors
 Author-email: CAI4CAI research group <contact@cai4cai.uk>
 License-Expression: Apache-2.0

{torchsparsegradutils-0.2.5 → torchsparsegradutils-0.2.6}/docs/source/conf.py

@@ -15,8 +15,8 @@ sys.path.insert(0, os.path.abspath("../../"))
 project = "torchsparsegradutils"
 copyright = "2026, CAI4CAI research group"
 author = "CAI4CAI research group"
-release = "0.2.5"
-version = "0.2.5"
+release = "0.2.6"
+version = "0.2.6"
 
 # -- General configuration ---------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration

{torchsparsegradutils-0.2.5 → torchsparsegradutils-0.2.6}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "torchsparsegradutils"
-version = "0.2.5"
+version = "0.2.6"
 description = "A collection of utility functions to work with PyTorch sparse tensors"
 readme = "README.md"
 requires-python = ">=3.10"

torchsparsegradutils-0.2.6/torchsparsegradutils/benchmarks/linear_cg_convergence.py

@@ -0,0 +1,224 @@
+#!/usr/bin/env python3
+"""Compare linear-CG convergence before and after the diagnostics fix.
+
+This benchmark focuses on numerical behavior rather than throughput. It loads
+the historical implementation from Git so both solvers run in one Python
+environment on identical deterministic inputs.
+
+Example
+-------
+python -m torchsparsegradutils.benchmarks.linear_cg_convergence \
+    --baseline-ref ea7b8f0 --repeats 100
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import statistics
+import subprocess
+import time
+import types
+import warnings
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Callable
+
+import torch
+
+from torchsparsegradutils.utils import linear_cg
+
+REPOSITORY_ROOT = Path(__file__).resolve().parents[2]
+LINEAR_CG_PATH = "torchsparsegradutils/utils/linear_cg.py"
+
+
+@dataclass(frozen=True)
+class Problem:
+    name: str
+    matrix: torch.Tensor
+    rhs: torch.Tensor
+    tolerance: float
+    max_iter: int
+    initial_guess: torch.Tensor | None = None
+
+
+class CountedMatmul:
+    def __init__(self, matrix: torch.Tensor):
+        self.matrix = matrix
+        self.calls = 0
+
+    def __call__(self, value: torch.Tensor) -> torch.Tensor:
+        self.calls += 1
+        return self.matrix @ value
+
+
+def load_historical_linear_cg(revision: str) -> Callable[..., torch.Tensor]:
+    """Load ``linear_cg`` from a Git revision without changing the worktree."""
+    command = ["git", "-C", str(REPOSITORY_ROOT), "show", f"{revision}:{LINEAR_CG_PATH}"]
+    try:
+        completed = subprocess.run(command, check=True, capture_output=True, text=True)
+    except FileNotFoundError as error:
+        raise RuntimeError("Cannot load the historical solver because the Git executable was not found") from error
+    except subprocess.CalledProcessError as error:
+        detail = error.stderr.strip() or "Git could not resolve the requested revision and file"
+        raise RuntimeError(
+            f"Cannot load the historical solver from revision {revision!r}. "
+            f"Run this benchmark from a Git checkout containing that revision. Git reported: {detail}"
+        ) from error
+    module = types.ModuleType("historical_linear_cg")
+    exec(compile(completed.stdout, f"{revision}:{LINEAR_CG_PATH}", "exec"), module.__dict__)
+    return module.linear_cg
+
+
+def make_problems(device: torch.device) -> list[Problem]:
+    dtype = torch.float64
+
+    size = 36
+    diagonal = torch.full((size,), 2.0, dtype=dtype, device=device)
+    off_diagonal = torch.full((size - 1,), -0.25, dtype=dtype, device=device)
+    tight_matrix = torch.diag(diagonal) + torch.diag(off_diagonal, diagonal=1) + torch.diag(off_diagonal, diagonal=-1)
+    tight_rhs = torch.zeros(size, dtype=dtype, device=device)
+    tight_rhs[size // 2] = 1
+
+    size = 40
+    multiple_matrix = torch.diag(torch.linspace(1.0, 100.0, size, dtype=dtype, device=device))
+    multiple_rhs = torch.zeros((size, 101), dtype=dtype, device=device)
+    multiple_rhs[:, -1] = 1
+
+    zero_matrix = torch.diag(torch.tensor([1.0, 2.0, 3.0], dtype=dtype, device=device))
+    zero_rhs = torch.zeros(3, dtype=dtype, device=device)
+
+    return [
+        Problem("tight_tolerance", tight_matrix, tight_rhs, tolerance=1e-12, max_iter=200),
+        Problem("multiple_rhs", multiple_matrix, multiple_rhs, tolerance=1e-4, max_iter=size),
+        Problem(
+            "zero_rhs_nonzero_guess",
+            zero_matrix,
+            zero_rhs,
+            tolerance=1e-5,
+            max_iter=20,
+            initial_guess=torch.ones_like(zero_rhs),
+        ),
+    ]
+
+
+def synchronize(device: torch.device) -> None:
+    if device.type == "cuda":
+        torch.cuda.synchronize(device)
+
+
+def relative_residual_per_rhs(matrix: torch.Tensor, solution: torch.Tensor, rhs: torch.Tensor) -> torch.Tensor:
+    if rhs.ndim == 1:
+        rhs = rhs.unsqueeze(-1)
+        solution = solution.unsqueeze(-1)
+    residual_norm = torch.linalg.vector_norm(rhs - matrix @ solution, dim=-2)
+    rhs_norm = torch.linalg.vector_norm(rhs, dim=-2)
+    return residual_norm / rhs_norm.masked_fill(rhs_norm.eq(0), 1)
+
+
+def run_once(
+    solver: Callable[..., torch.Tensor],
+    problem: Problem,
+    *,
+    fixed: bool,
+) -> tuple[torch.Tensor, int, int, str]:
+    matmul = CountedMatmul(problem.matrix)
+    arguments = {
+        "tolerance": problem.tolerance,
+        "max_iter": problem.max_iter,
+        "initial_guess": problem.initial_guess,
+    }
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", UserWarning)
+        if fixed:
+            solution, info = solver(matmul, problem.rhs, return_info=True, **arguments)
+            return solution, info.iterations, matmul.calls, info.reason
+        solution = solver(matmul, problem.rhs, **arguments)
+    # The historical solver performs one initial matvec followed by one per
+    # iteration and does not recompute the true residual before returning.
+    return solution, max(0, matmul.calls - 1), matmul.calls, "not_reported"
+
+
+def measure(
+    solver: Callable[..., torch.Tensor],
+    problem: Problem,
+    *,
+    fixed: bool,
+    warmup: int,
+    repeats: int,
+) -> dict[str, object]:
+    for _ in range(warmup):
+        run_once(solver, problem, fixed=fixed)
+    synchronize(problem.rhs.device)
+
+    durations = []
+    for _ in range(repeats):
+        start = time.perf_counter()
+        run_once(solver, problem, fixed=fixed)
+        synchronize(problem.rhs.device)
+        durations.append(time.perf_counter() - start)
+
+    solution, iterations, matvecs, reason = run_once(solver, problem, fixed=fixed)
+    true_residual = relative_residual_per_rhs(problem.matrix, solution, problem.rhs)
+    reference = torch.linalg.solve(problem.matrix, problem.rhs)
+    error = torch.linalg.vector_norm(solution - reference) / torch.linalg.vector_norm(reference).clamp_min(1)
+
+    return {
+        "iterations": iterations,
+        "matvecs": matvecs,
+        "reason": reason,
+        "true_relative_residual_max": float(true_residual.max()),
+        "unconverged_rhs": int((true_residual > problem.tolerance).sum()),
+        "relative_solution_error": float(error),
+        "median_time_ms": statistics.median(durations) * 1e3,
+    }
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--baseline-ref", default="ea7b8f0", help="Git revision containing the historical solver")
+    parser.add_argument("--device", default="cpu", help="PyTorch device (default: cpu)")
+    parser.add_argument("--warmup", type=int, default=5)
+    parser.add_argument("--repeats", type=int, default=50)
+    arguments = parser.parse_args()
+    if arguments.warmup < 0 or arguments.repeats < 1:
+        parser.error("--warmup must be nonnegative and --repeats must be positive")
+    return arguments
+
+
+def main() -> None:
+    arguments = parse_args()
+    device = torch.device(arguments.device)
+    historical_linear_cg = load_historical_linear_cg(arguments.baseline_ref)
+    results = {}
+    for problem in make_problems(device):
+        results[problem.name] = {
+            "baseline": measure(
+                historical_linear_cg,
+                problem,
+                fixed=False,
+                warmup=arguments.warmup,
+                repeats=arguments.repeats,
+            ),
+            "fixed": measure(
+                linear_cg,
+                problem,
+                fixed=True,
+                warmup=arguments.warmup,
+                repeats=arguments.repeats,
+            ),
+        }
+
+    payload = {
+        "baseline_ref": arguments.baseline_ref,
+        "device": str(device),
+        "torch_version": torch.__version__,
+        "warmup": arguments.warmup,
+        "repeats": arguments.repeats,
+        "results": results,
+    }
+    print(json.dumps(payload, indent=2, sort_keys=True))
+
+
+if __name__ == "__main__":
+    main()

torchsparsegradutils-0.2.6/torchsparsegradutils/tests/test_linear_cg.py

@@ -0,0 +1,366 @@
+import pytest
+import torch
+
+from torchsparsegradutils.utils.linear_cg import linear_cg
+
+
+# Test basic CG solve for vectors and matrices
+def test_cg():
+    size = 100
+    # SPD matrix
+    matrix = torch.randn(size, size, dtype=torch.float64)
+    matrix = matrix.matmul(matrix.mT)
+    matrix.div_(torch.linalg.vector_norm(matrix)).add_(torch.eye(size, dtype=torch.float64) * 1e-1)
+    # single RHS
+    rhs = torch.randn(size, dtype=torch.float64)
+    solves = linear_cg(matrix.matmul, rhs=rhs, max_iter=size)
+    init = torch.randn(size, dtype=torch.float64)
+    solves_init = linear_cg(matrix.matmul, rhs=rhs, max_iter=size, initial_guess=init)
+    # reference solve
+    chol = torch.linalg.cholesky(matrix)
+    actual = torch.cholesky_solve(rhs.unsqueeze(1), chol).squeeze()
+    assert torch.allclose(solves, actual, atol=1e-3, rtol=1e-4)
+    assert torch.allclose(solves_init, actual, atol=1e-3, rtol=1e-4)
+    # multiple RHS
+    rhs_mat = torch.randn(size, 50, dtype=torch.float64)
+    solves = linear_cg(matrix.matmul, rhs=rhs_mat, max_iter=size)
+    init_mat = torch.randn(size, 50, dtype=torch.float64)
+    solves_init = linear_cg(matrix.matmul, rhs=rhs_mat, max_iter=size, initial_guess=init_mat)
+    actual_mat = torch.cholesky_solve(rhs_mat, chol)
+    assert torch.allclose(solves, actual_mat, atol=1e-3, rtol=1e-4)
+    assert torch.allclose(solves_init, actual_mat, atol=1e-3, rtol=1e-4)
+
+
+# Test CG with tridiagonal outputs
+def test_cg_with_tridiag():
+    size = 10
+    matrix = torch.randn(size, size, dtype=torch.float64)
+    matrix = matrix.matmul(matrix.mT)
+    matrix.div_(torch.linalg.vector_norm(matrix)).add_(torch.eye(size, dtype=torch.float64) * 1e-1)
+    rhs = torch.randn(size, 50, dtype=torch.float64)
+    solves, t_mats = linear_cg(
+        matrix.matmul, rhs=rhs, n_tridiag=5, max_tridiag_iter=10, max_iter=size, tolerance=0, eps=1e-15
+    )
+    chol = torch.linalg.cholesky(matrix)
+    actual = torch.cholesky_solve(rhs, chol)
+    assert torch.allclose(solves, actual, atol=1e-3, rtol=1e-4)
+    eigs = torch.linalg.eigvalsh(matrix)
+    for i in range(5):
+        approx = torch.linalg.eigvalsh(t_mats[i])
+        assert torch.allclose(eigs, approx, atol=1e-3, rtol=1e-4)
+
+
+# Device parameterized CG tests
+@pytest.mark.parametrize("batch", [None, 5])
+def test_batch_cg(batch):
+    size = 100
+    shape = (batch, size, size) if batch else (size, size)
+    matrix = torch.randn(*shape, dtype=torch.float64)
+    matrix = matrix.matmul(matrix.mT)
+    matrix.div_(torch.linalg.vector_norm(matrix)).add_(torch.eye(size, dtype=torch.float64) * 1e-1)
+    b_shape = (batch, size, 50) if batch else (size, 50)
+    rhs = torch.randn(*b_shape, dtype=torch.float64)
+    solves = linear_cg(matrix.matmul, rhs=rhs, max_iter=size)
+    chol = torch.linalg.cholesky(matrix)
+    actual = torch.cholesky_solve(rhs, chol)
+    assert torch.allclose(solves, actual, atol=1e-3, rtol=1e-4)
+
+
+@pytest.mark.parametrize("batch", [None, 5])
+def test_batch_cg_with_tridiag(batch):
+    size = 10
+    shape = (batch, size, size) if batch else (size, size)
+    matrix = torch.randn(*shape, dtype=torch.float64)
+    matrix = matrix.matmul(matrix.mT)
+    matrix.div_(torch.linalg.vector_norm(matrix)).add_(torch.eye(size, dtype=torch.float64) * 1e-1)
+    b_shape = (batch, size, 10) if batch else (size, 10)
+    rhs = torch.randn(*b_shape, dtype=torch.float64)
+    solves, t_mats = linear_cg(
+        matrix.matmul, rhs=rhs, n_tridiag=8, max_iter=size, max_tridiag_iter=10, tolerance=0, eps=1e-30
+    )
+    chol = torch.linalg.cholesky(matrix)
+    actual = torch.cholesky_solve(rhs, chol)
+    assert torch.allclose(solves, actual, atol=1e-3, rtol=1e-4)
+    batch_dim = 5 if batch else 1
+    for i in range(batch_dim):
+        eigs = torch.linalg.eigvalsh(matrix[i] if batch else matrix)
+        for j in range(8):
+            approx = torch.linalg.eigvalsh(t_mats[j, i] if batch else t_mats[j])
+            assert torch.allclose(eigs, approx, atol=1e-3, rtol=1e-4)
+
+
+# Test CG initialization reuse
+def test_batch_cg_init():
+    batch = 5
+    size = 100
+    matrix = torch.randn(batch, size, size, dtype=torch.float64)
+    matrix = matrix.matmul(matrix.mT)
+    matrix.div_(torch.linalg.vector_norm(matrix)).add_(torch.eye(size, dtype=torch.float64) * 1e-1)
+    rhs = torch.randn(batch, size, 50, dtype=torch.float64)
+    solves = linear_cg(matrix.matmul, rhs=rhs, max_iter=size, max_tridiag_iter=0)
+    solves_init = linear_cg(matrix.matmul, rhs=rhs, max_iter=1, initial_guess=solves, max_tridiag_iter=0)
+    chol = torch.linalg.cholesky(matrix)
+    actual = torch.cholesky_solve(rhs, chol)
+    assert torch.allclose(solves_init, actual, atol=1e-3, rtol=1e-4)
+
+
+def test_tight_tolerance_does_not_stall_at_historical_absolute_epsilon():
+    size = 36
+    diagonal = torch.full((size,), 2.0, dtype=torch.float64)
+    off_diagonal = torch.full((size - 1,), -0.25, dtype=torch.float64)
+    matrix = torch.diag(diagonal) + torch.diag(off_diagonal, 1) + torch.diag(off_diagonal, -1)
+    rhs = torch.zeros(size, dtype=torch.float64)
+    rhs[size // 2] = 1
+
+    solution, info = linear_cg(matrix, rhs, tolerance=1e-12, max_iter=200, return_info=True)
+    expected = torch.linalg.solve(matrix, rhs)
+
+    torch.testing.assert_close(solution, expected, rtol=1e-10, atol=1e-12)
+    assert info.reason == "converged"
+    assert info.converged.all()
+    assert info.true_relative_residual.max() <= 1e-12
+    assert info.iterations <= size
+
+
+def test_all_rhs_convergence_does_not_hide_one_hard_column():
+    size = 40
+    matrix = torch.diag(torch.linspace(1.0, 100.0, size, dtype=torch.float64))
+    rhs = torch.zeros((size, 101), dtype=torch.float64)
+    rhs[:, -1] = 1
+
+    with pytest.warns(UserWarning):
+        _, mean_info = linear_cg(
+            matrix,
+            rhs,
+            tolerance=1e-4,
+            max_iter=size,
+            convergence_reduction="mean",
+            return_info=True,
+        )
+    all_solution, all_info = linear_cg(
+        matrix,
+        rhs,
+        tolerance=1e-4,
+        max_iter=size,
+        convergence_reduction="all",
+        return_info=True,
+    )
+
+    assert mean_info.reason == "mean_converged"
+    assert not mean_info.converged[..., -1].all()
+    assert all_info.reason == "converged"
+    assert all_info.converged.all()
+    assert all_info.true_relative_residual.max() <= 1e-4
+    torch.testing.assert_close(all_solution, torch.linalg.solve(matrix, rhs), rtol=2e-4, atol=1e-6)
+
+
+def test_zero_rhs_returns_zero_even_with_nonzero_initial_guess():
+    matrix = torch.diag(torch.tensor([1.0, 2.0, 3.0], dtype=torch.float64))
+    rhs = torch.zeros(3, dtype=torch.float64)
+    initial_guess = torch.ones(3, dtype=torch.float64)
+
+    solution, info = linear_cg(matrix, rhs, initial_guess=initial_guess, tolerance=0, return_info=True)
+
+    torch.testing.assert_close(solution, torch.zeros_like(rhs), rtol=0, atol=0)
+    assert info.iterations == 0
+    assert info.converged.all()
+    assert info.true_relative_residual.max() == 0
+    assert info.tolerance == 0
+
+
+def test_mixed_zero_and_nonzero_rhs_columns_are_supported():
+    matrix = torch.diag(torch.tensor([1.0, 2.0, 4.0], dtype=torch.float64))
+    rhs = torch.stack((torch.zeros(3, dtype=torch.float64), torch.ones(3, dtype=torch.float64)), dim=-1)
+
+    solution, info = linear_cg(matrix, rhs, tolerance=1e-12, return_info=True)
+
+    expected = torch.linalg.solve(matrix, rhs)
+    torch.testing.assert_close(solution, expected, rtol=1e-12, atol=1e-12)
+    torch.testing.assert_close(solution[:, 0], torch.zeros(3, dtype=torch.float64), rtol=0, atol=0)
+    assert info.converged.all()
+
+
+def test_preconditioner_and_breakdown_validation():
+    diagonal = torch.tensor([1.0, 2.0, 4.0], dtype=torch.float64)
+    matrix = torch.diag(diagonal)
+    rhs = torch.ones(3, dtype=torch.float64)
+
+    solution, info = linear_cg(
+        matrix,
+        rhs,
+        preconditioner=lambda value: value / diagonal.unsqueeze(-1),
+        return_info=True,
+    )
+    torch.testing.assert_close(solution, torch.linalg.solve(matrix, rhs), rtol=1e-12, atol=1e-12)
+    assert info.reason == "converged"
+
+    indefinite = torch.diag(torch.tensor([1.0, -1.0, 2.0], dtype=torch.float64))
+    with pytest.raises(RuntimeError, match=r"p\^T A p"):
+        linear_cg(indefinite, torch.tensor([0.0, 1.0, 0.0], dtype=torch.float64))
+
+    with pytest.raises(RuntimeError, match="preconditioned residual inner product"):
+        linear_cg(matrix, rhs, preconditioner=lambda value: -value)
+
+
+def test_user_eps_raises_instead_of_freezing_active_column():
+    matrix = torch.eye(3, dtype=torch.float64)
+    rhs = torch.ones(3, dtype=torch.float64)
+
+    with pytest.raises(RuntimeError, match="at least eps"):
+        linear_cg(
+            matrix,
+            rhs,
+            initial_guess=(1 - 1e-6) * rhs,
+            tolerance=1e-8,
+            eps=1e-10,
+            max_tridiag_iter=0,
+        )
+
+
+def test_stopped_initial_updates_warn_and_report_reason():
+    matrix = torch.eye(3, dtype=torch.float64)
+    rhs = torch.ones(3, dtype=torch.float64)
+
+    with pytest.warns(UserWarning, match="did not converge"):
+        _, info = linear_cg(
+            matrix,
+            rhs,
+            initial_guess=0.5 * rhs,
+            tolerance=0.1,
+            stop_updating_after=0.6,
+            return_info=True,
+        )
+
+    assert info.iterations == 0
+    assert info.reason == "stopped_updating"
+    assert info.tolerance == 0.1
+    assert not info.converged.any()
+
+
+def test_recursive_convergence_is_distinguished_from_true_convergence():
+    calls = 0
+
+    def matmul_with_final_residual_drift(value):
+        nonlocal calls
+        calls += 1
+        return 1.1 * value if calls == 3 else value
+
+    with pytest.warns(UserWarning, match="did not converge"):
+        _, info = linear_cg(
+            matmul_with_final_residual_drift,
+            torch.ones(3, dtype=torch.float64),
+            tolerance=1e-12,
+            max_tridiag_iter=0,
+            return_info=True,
+        )
+
+    assert info.iterations == 1
+    assert info.reason == "recursive_converged"
+    assert info.recursive_relative_residual.max() == 0
+    assert info.true_relative_residual.max() > info.tolerance
+
+
+def test_tridiagonalization_can_return_info():
+    matrix = torch.diag(torch.tensor([1.0, 2.0, 3.0], dtype=torch.float64))
+    rhs = torch.ones(3, dtype=torch.float64)
+
+    solution, tridiagonal, info = linear_cg(
+        matrix,
+        rhs,
+        n_tridiag=1,
+        tolerance=0,
+        max_iter=3,
+        max_tridiag_iter=3,
+        return_info=True,
+    )
+
+    torch.testing.assert_close(solution, torch.linalg.solve(matrix, rhs), rtol=1e-12, atol=1e-12)
+    assert tridiagonal.shape == (1, 3, 3)
+    assert info.iterations == 3
+    assert info.matvecs == 5
+
+
+@pytest.mark.parametrize(
+    ("kwargs", "match"),
+    [
+        ({"tolerance": float("nan")}, "tolerance"),
+        ({"tolerance": float("inf")}, "tolerance"),
+        ({"tolerance": -1.0}, "tolerance"),
+        ({"eps": float("nan")}, "eps"),
+        ({"eps": float("inf")}, "eps"),
+        ({"eps": 0.0}, "eps"),
+        ({"stop_updating_after": float("nan")}, "stop_updating_after"),
+        ({"stop_updating_after": float("inf")}, "stop_updating_after"),
+        ({"stop_updating_after": -1.0}, "stop_updating_after"),
+        ({"convergence_reduction": "median"}, "convergence_reduction"),
+        ({"min_iter": -1}, "min_iter"),
+    ],
+)
+def test_invalid_solver_settings_raise(kwargs, match):
+    with pytest.raises(ValueError, match=match):
+        linear_cg(torch.eye(2, dtype=torch.float64), torch.ones(2, dtype=torch.float64), **kwargs)
+
+
+def test_eps_must_be_representable_in_rhs_dtype():
+    with pytest.raises(ValueError, match="representable"):
+        linear_cg(torch.eye(2), torch.ones(2), eps=1e-100)
+
+
+def test_rhs_and_operator_output_validation():
+    matrix = torch.eye(2, dtype=torch.float64)
+    rhs = torch.ones(2, dtype=torch.float64)
+
+    with pytest.raises(TypeError, match="floating-point"):
+        linear_cg(torch.eye(2, dtype=torch.int64), torch.ones(2, dtype=torch.int64))
+    with pytest.raises(RuntimeError, match="matmul_closure output"):
+        linear_cg(lambda value: value[:-1], rhs)
+    with pytest.raises(RuntimeError, match="preconditioner output"):
+        linear_cg(matrix, rhs, preconditioner=lambda value: value.to(torch.float32))
+
+
+def test_vector_rank_is_preserved_with_vector_initial_guess():
+    matrix = torch.diag(torch.tensor([1.0, 2.0, 3.0], dtype=torch.float64))
+    rhs = torch.ones(3, dtype=torch.float64)
+    solution = linear_cg(matrix, rhs, initial_guess=torch.zeros_like(rhs))
+    assert solution.shape == rhs.shape
+
+    column_guess_solution = linear_cg(matrix, rhs, initial_guess=torch.zeros((3, 1), dtype=torch.float64))
+    assert column_guess_solution.shape == rhs.shape
+
+    with pytest.raises(ValueError, match="initial_guess must have shape"):
+        linear_cg(matrix, rhs, initial_guess=torch.zeros((3, 2), dtype=torch.float64))
+
+
+def test_min_iter_is_honored_for_converged_initial_guess():
+    matrix = torch.eye(3, dtype=torch.float64)
+    rhs = torch.ones(3, dtype=torch.float64)
+
+    solution, info = linear_cg(
+        matrix,
+        rhs,
+        initial_guess=rhs,
+        min_iter=3,
+        max_iter=5,
+        max_tridiag_iter=0,
+        return_info=True,
+    )
+
+    torch.testing.assert_close(solution, rhs, rtol=0, atol=0)
+    assert info.iterations == 3
+    assert info.matvecs == 5
+    assert info.reason == "converged"
+    assert info.converged.all()
+
+
+def test_multiple_batch_dimensions_are_supported():
+    with pytest.raises(ValueError, match="at least one dimension"):
+        linear_cg(torch.ones((1, 1), dtype=torch.float64), torch.tensor(1.0, dtype=torch.float64))
+
+    matrix = torch.diag(torch.tensor([1.0, 2.0, 3.0], dtype=torch.float64))
+    rhs = torch.ones((2, 4, 3, 2), dtype=torch.float64)
+
+    solution = linear_cg(matrix, rhs, tolerance=1e-12)
+    expected = torch.linalg.solve(matrix, rhs)
+
+    assert solution.shape == rhs.shape
+    torch.testing.assert_close(solution, expected, rtol=1e-12, atol=1e-12)