nlsq.svd_fallback module

SVD computation with GPU/CPU fallback for robustness.

This module provides: - compute_svd_with_fallback: Deterministic full SVD with GPU/CPU/NumPy fallback chain

IMPORTANT: As of v0.3.5, this module uses ONLY full deterministic SVD. Randomized/approximate SVD has been completely removed because it causes optimization divergence in iterative least-squares solvers.

Historical note (v0.3.1-v0.3.4):

Randomized SVD was available but caused 3-25x worse fitting errors in iterative least-squares applications due to accumulated approximation error across trust-region iterations. See tests/test_svd_regression.py for evidence.

nlsq.stability.svd_fallback.is_gpu_error(error)[source]

Check if an exception indicates a GPU/CUDA-specific failure.

Prefers type-based matching against jaxlib.xla_extension.XlaRuntimeError, which is stable across JAX versions. Falls back to string heuristics for older JAX (<0.8) or error types that don’t inherit from XlaRuntimeError: - Legacy: “cuSolver internal error” (JAX <0.8) - FFI: “No FFI handler registered for cusolver_gesvdj_ffi” (JAX >=0.8) - XLA status: “INTERNAL: …” (gRPC/XLA status code prefix, case-sensitive)

nlsq.stability.svd_fallback.compute_svd_with_fallback(J_h, full_matrices=False)[source]

Compute full deterministic SVD with multiple fallback strategies.

This is the primary SVD function for NLSQ. It uses full (exact) SVD to ensure numerical precision and reproducibility in optimization.

Fallback chain: 1. JAX GPU SVD (if GPU available) 2. JAX CPU SVD (if GPU fails with cuSolver error) 3. NumPy SVD (last resort)

Parameters:

  • J_h (jnp.ndarray) – Jacobian matrix in hat space

  • full_matrices (bool) – Whether to compute full matrices (default: False for efficiency)

Returns:

  • U (jnp.ndarray) – Left singular vectors

  • s (jnp.ndarray) – Singular values (sorted in descending order)

  • V (jnp.ndarray) – Right singular vectors (note: V is transposed back, NOT Vt)

Return type:

tuple[Array, Array, Array]

nlsq.stability.svd_fallback.initialize_gpu_safely()[source]

Initialize GPU with proper memory settings to avoid cuSolver issues.

Overview

The svd_fallback module provides full deterministic SVD computation with GPU/CPU/NumPy fallback chain for robust numerical precision.

Changed in version 0.3.5: Randomized SVD has been completely removed. All SVD operations now use full deterministic SVD for numerical precision and reproducibility in optimization.

Key Features

  • Full deterministic SVD for numerical precision

  • GPU/CPU fallback for robust SVD computation

  • NumPy fallback as last resort if JAX fails

  • Reproducible results across runs and platforms

Functions

compute_svd_with_fallback(J_h[, full_matrices])

Compute full deterministic SVD with multiple fallback strategies.

compute_svd_with_fallback (Primary API)

The recommended function for SVD in NLSQ:

from nlsq.stability.svd_fallback import compute_svd_with_fallback
import jax.numpy as jnp

# Matrix of any size
A = jnp.ones((100_000, 50))

# Full deterministic SVD with automatic fallback
U, s, V = compute_svd_with_fallback(A, full_matrices=False)

print(f"U shape: {U.shape}")  # (100000, 50)
print(f"s shape: {s.shape}")  # (50,)
print(f"V shape: {V.shape}")  # (50, 50)

Fallback Sequence:

  1. JAX GPU SVD (jax.scipy.linalg.svd on GPU)

  2. JAX CPU SVD (automatic fallback if GPU fails with cuSolver error)

  3. NumPy SVD (last resort if JAX CPU also fails)

GPU/CPU Fallback

Handle GPU failures gracefully:

from nlsq.stability.svd_fallback import compute_svd_with_fallback
import jax.numpy as jnp

# Matrix that might cause numerical issues
A = jnp.array([[1e10, 1.0], [1.0, 1e-10]])

# SVD with automatic GPU→CPU→NumPy fallback
U, s, V = compute_svd_with_fallback(A, full_matrices=False)

print(f"Singular values: {s}")

Why Full SVD Only?

Randomized SVD (available in v0.3.1-v0.3.4) was removed because it caused optimization divergence in iterative least-squares solvers:

  • Approximation error accumulates across trust-region iterations

  • Early termination at worse local minima

  • 3-25x worse fitting errors in sensitive applications

Evidence from XPCS fitting (50K points, 13 params):

SVD Method

D0 Error

Alpha Error

Iterations

Full SVD (v0.3.0)

9.74%

0.59%

15

Randomized SVD

30.18%

14.66%

6

See tests/test_svd_regression.py for detailed regression tests.

See Also