Troubleshooting Guide

This guide helps you diagnose and fix common issues when using NLSQ.

---

Installation Issues

Issue: ModuleNotFoundError: No module named 'nlsq'

Cause: NLSQ not installed

Solution:

pip install nlsq

# Or for development:
git clone https://github.com/imewei/NLSQ.git
cd NLSQ
pip install -e .

Verify installation:

import nlsq

print(nlsq.__version__)

Issue: ImportError: JAX requires NumPy >= 1.21

Cause: Incompatible NumPy version

Solution:

pip install --upgrade numpy>=1.21
pip install --upgrade jax jaxlib

Issue: CUDA version mismatch

Error:

RuntimeError: jaxlib version 0.4.1 is newer than and incompatible with jax version 0.3.25

Solution:

# Uninstall and reinstall with matching versions
pip uninstall jax jaxlib
pip install "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

# Or for CPU only:
pip install "jax[cpu]"

Check CUDA version:

nvidia-smi  # Look for CUDA Version
nvcc --version

---

GPU/TPU Issues

Issue: No GPU detected (using CPU instead)

Symptoms:

import jax

print(jax.devices())
# Output: [CpuDevice(id=0)]  # Should be GpuDevice

Solutions:

1. Check GPU availability:

nvidia-smi  # Should show GPU info

2. Reinstall JAX with CUDA support:

# For CUDA 12.x
pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

# For CUDA 11.x
pip install --upgrade "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

3. Check CUDA environment variables:

echo $CUDA_HOME
echo $LD_LIBRARY_PATH

4. Verify JAX can see CUDA:

import jax

print(jax.local_devices())  # Should include GPU

Issue: GPU out of memory (OOM)

Error:

RuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 1234567890 bytes

Solutions:

1. Limit memory preallocation:

import os

os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = "0.75"  # Use 75% of GPU

from nlsq import curve_fit

2. Use chunking for large datasets:

from nlsq.streaming.large_dataset import fit_large_dataset

popt, pcov, info = fit_large_dataset(
    model,
    x,
    y,
    p0=[2, 1],
    memory_limit_gb=4.0,  # Limit GPU memory usage
    chunk_size=100_000,
)

3. Use memory-efficient solver:

popt, pcov = curve_fit(
    model, x, y, p0=[2, 1], solver="cg"  # More memory efficient than 'svd'
)

4. Clear GPU cache:

import jax

jax.clear_backends()  # Release all GPU memory

Issue: GPU slower than CPU

Cause: Dataset too small (JIT overhead dominates)

Solution:

import os

# For datasets < 10K points, use CPU
if len(x) < 10000:
    os.environ["JAX_PLATFORM_NAME"] = "cpu"

from nlsq import curve_fit

Or benchmark both:

import time

# CPU timing
os.environ["JAX_PLATFORM_NAME"] = "cpu"
start = time.time()
popt_cpu, _ = curve_fit(model, x, y, p0=[2, 1])
cpu_time = time.time() - start

# GPU timing
os.environ["JAX_PLATFORM_NAME"] = "gpu"
start = time.time()
popt_gpu, _ = curve_fit(model, x, y, p0=[2, 1])
gpu_time = time.time() - start

print(f"CPU: {cpu_time:.3f}s, GPU: {gpu_time:.3f}s")

---

Convergence Problems

Issue: RuntimeError: Optimal parameters not found

Cause: Optimization failed to converge

Diagnosis:

try:
    popt, pcov = curve_fit(model, x, y, p0=[2, 1])
except RuntimeError as e:
    print(f"Error: {e}")
    # Get more details
    result = curve_fit(model, x, y, p0=[2, 1], full_output=True)
    print(f"Status: {result.status}")
    print(f"Message: {result.message}")

Solutions:

1. Improve initial guess (``p0``):

# Bad: p0 far from solution
popt, pcov = curve_fit(model, x, y, p0=[100, 0.001])  # May fail

# Good: p0 closer to expected values
popt, pcov = curve_fit(model, x, y, p0=[2, 1])  # More likely to succeed

2. Set realistic bounds:

popt, pcov = curve_fit(
    model, x, y, p0=[2, 1], bounds=([0, 0], [10, 5])  # Constrain search space
)

3. Increase tolerance:

popt, pcov = curve_fit(
    model,
    x,
    y,
    p0=[2, 1],
    ftol=1e-6,  # Default: 1e-8 (looser tolerance)
    xtol=1e-6,
    gtol=1e-6,
)

4. Increase max iterations:

popt, pcov = curve_fit(
    model, x, y, p0=[2, 1], max_nfev=10000  # Default: 100 * (n_params + 1)
)

5. Scale your data:

# Bad: x in [0, 1e6], y in [1e-10, 1e-8]
popt, pcov = curve_fit(model, x, y, p0=[2, 1])  # May fail

# Good: Scale to reasonable ranges
x_scaled = x / 1e6  # Now in [0, 1]
y_scaled = y * 1e10  # Now in [1, 100]
popt, pcov = curve_fit(model, x_scaled, y_scaled, p0=[2, 1])

# Unscale results
popt[0] = popt[0] / 1e10  # Unscale amplitude parameter

Issue: Fit converges but results are wrong

Cause: Local minimum or poor initial guess

Solutions:

1. Try multiple initial guesses:

import numpy as np

p0_guesses = [[1, 0.5], [2, 1.0], [5, 2.0], [10, 0.1]]

best_cost = np.inf
best_popt = None

for p0 in p0_guesses:
    try:
        popt, pcov = curve_fit(model, x, y, p0=p0)
        cost = np.sum((y - model(x, *popt)) ** 2)
        if cost < best_cost:
            best_cost = cost
            best_popt = popt
    except RuntimeError:
        continue

print(f"Best fit: {best_popt}, cost: {best_cost}")

2. Visualize the fit:

import matplotlib.pyplot as plt

popt, pcov = curve_fit(model, x, y, p0=[2, 1])

plt.figure(figsize=(10, 4))

# Plot 1: Data and fit
plt.subplot(1, 2, 1)
plt.plot(x, y, "o", label="Data")
plt.plot(x, model(x, *popt), "-", label="Fit")
plt.legend()

# Plot 2: Residuals
plt.subplot(1, 2, 2)
residuals = y - model(x, *popt)
plt.plot(x, residuals, "o")
plt.axhline(0, color="r", linestyle="--")
plt.ylabel("Residuals")

plt.tight_layout()
plt.show()

Issue: Covariance matrix has inf or nan

Cause: Jacobian is singular or near-singular

Solutions:

1. Check parameter identifiability:

# Some parameters may not be identifiable from data
def model(x, a, b, c):
    return a * jnp.exp(-b * x) + c


# If data doesn't cover x=0, 'a+c' not separately identifiable

# Solution: Fix one parameter or add constraints
popt, pcov = curve_fit(lambda x, a, b: model(x, a, b, c=0.5), x, y, p0=[2, 1])  # Fix c

2. Add regularization via bounds:

popt, pcov = curve_fit(
    model,
    x,
    y,
    p0=[2, 1, 0.5],
    bounds=([0, 0, 0], [10, 5, 2]),  # Prevent singular solutions
)

3. Check for redundant parameters:

# Bad: Parameters are correlated
def model(x, a, b, c, d):
    return a * jnp.exp(-b * x) + c * jnp.exp(-d * x)


# If b ≈ d, parameters are redundant


# Good: Use fewer parameters
def model(x, a, b):
    return a * jnp.exp(-b * x)

---

Performance Issues

Issue: First fit is very slow

Cause: JIT compilation overhead

Solution: This is expected. Subsequent fits will be much faster.

from nlsq import CurveFit

# Create reusable fitter
fitter = CurveFit()

# First call: slow (compilation + execution)
popt1, pcov1 = fitter.curve_fit(model, x1, y1, p0=[2, 1])  # ~500ms

# Subsequent calls: fast (execution only)
popt2, pcov2 = fitter.curve_fit(model, x2, y2, p0=[2, 1])  # ~30ms
popt3, pcov3 = fitter.curve_fit(model, x3, y3, p0=[2, 1])  # ~30ms

Issue: Fit is slower than SciPy

Diagnosis:

1. Check dataset size:

print(f"Dataset size: {len(x)}")

If < 1000 points: NLSQ overhead may not be worth it. Use SciPy.

2. Check if GPU is being used:

import jax

print(f"Devices: {jax.devices()}")

3. Benchmark with timing:

popt, pcov, res, post_time, compile_time = curve_fit(
    model, x, y, p0=[2, 1], timeit=True
)

print(f"Compile time: {compile_time:.3f}s")
print(f"Execution time: {post_time:.3f}s")

Solutions:

• Use GPU for datasets > 10K points

• Use CurveFit class for multiple fits

• See Performance Optimization Guide

Issue: Memory usage keeps growing

Cause: JIT cache growing or memory not being released

Solutions:

1. Clear JIT cache periodically:

import jax

# After many fits
jax.clear_caches()

2. Disable JIT caching (not recommended):

from jax import config

config.update("jax_compilation_cache_dir", "")

3. Use chunking for large datasets:

from nlsq.streaming.large_dataset import fit_large_dataset

popt, pcov, info = fit_large_dataset(model, x_large, y_large, memory_limit_gb=4.0)

---

Memory Issues

Issue: MemoryError during fit

Solutions:

1. Use chunking:

from nlsq.streaming.large_dataset import fit_large_dataset

popt, pcov, info = fit_large_dataset(
    model,
    x,
    y,
    p0=[2, 1],
    chunk_size=100_000,  # Process 100K points at a time
    memory_limit_gb=4.0,
)

2. Use minibatch solver:

popt, pcov = curve_fit(model, x, y, p0=[2, 1], solver="minibatch", batch_size=50_000)

3. Use streaming optimizer for very large datasets:

from nlsq import curve_fit_large

popt, pcov = curve_fit_large(model, x, y, p0=p0)

---

Numerical Stability Issues

Issue: RuntimeError: NaN or Inf encountered

Causes: - Overflow/underflow in model function - Division by zero - Log of negative number

Solutions:

1. Add numerical safeguards:

import jax.numpy as jnp


# Bad: Can overflow or divide by zero
def model(x, a, b):
    return a / (1 + jnp.exp(-b * x))


# Good: Add safeguards
def model(x, a, b):
    # Clip to prevent overflow
    z = jnp.clip(-b * x, -100, 100)
    return a / (1 + jnp.exp(z))

2. Use stable numerical functions:

# Bad: log(exp(x)) can overflow
result = jnp.log(jnp.exp(x))

# Good: Use logsumexp
result = x  # Equivalent but stable

3. Check input data:

import numpy as np

# Check for inf/nan
assert np.all(np.isfinite(x))
assert np.all(np.isfinite(y))

# Check for very large/small values
print(f"x range: [{x.min()}, {x.max()}]")
print(f"y range: [{y.min()}, {y.max()}]")

Issue: Ill-conditioned Jacobian

Symptoms: - Large uncertainty estimates - Covariance matrix has very large or very small values - Warning: “Covariance cannot be estimated”

Solutions:

1. Scale parameters:

# Bad: Parameters have very different scales
def model(x, a, b):
    return a * jnp.exp(-b * x)  # a ~ 1e6, b ~ 1e-6


# Good: Rescale inside model
def model(x, a_scaled, b_scaled):
    a = a_scaled * 1e6
    b = b_scaled * 1e-6
    return a * jnp.exp(-b * x)


# Fit with scaled parameters
popt_scaled, pcov = curve_fit(model, x, y, p0=[1, 1])

# Unscale results
a_fit = popt_scaled[0] * 1e6
b_fit = popt_scaled[1] * 1e-6

2. Use parameter scaling:

popt, pcov = curve_fit(
    model, x, y, p0=[2, 1], x_scale="jac"  # Automatic parameter scaling
)

3. Check condition number:

from nlsq.diagnostics import check_condition_number

result = curve_fit(model, x, y, p0=[2, 1])
cond = check_condition_number(result.jac)

if cond > 1e10:
    print(f"Warning: Ill-conditioned (κ = {cond:.2e})")

---

API and Usage Errors

Issue: TypeError: curve_fit() got an unexpected keyword argument

Cause: Using SciPy-specific arguments in NLSQ

Solution:

# SciPy-only arguments (not supported in NLSQ):
# - full_output (use return_eval=True instead)
# - epsfcn, factor, diag (LM-specific)

# NLSQ equivalent:
popt, pcov = curve_fit(model, x, y, return_eval=False)  # Instead of full_output=False

Issue: ValueError: p0 must be a 1-D array

Cause: Incorrect p0 format

Solutions:

# Bad
popt, pcov = curve_fit(model, x, y, p0=[[2, 1]])  # 2D array

# Good
popt, pcov = curve_fit(model, x, y, p0=[2, 1])  # 1D array or list

Issue: ValueError: Residuals are not finite

Cause: Model returns inf/nan

Debug:

# Test model manually
p_test = [2, 1]
y_model = model(x, *p_test)
print(f"Model output finite: {np.all(np.isfinite(y_model))}")

# Check for specific issues
print(f"Contains NaN: {np.any(np.isnan(y_model))}")
print(f"Contains Inf: {np.any(np.isinf(y_model))}")

---

JAX-Specific Issues

Issue: TypeError: jax.numpy function called with non-jax array

Cause: Mixing NumPy and JAX arrays incorrectly

Solution:

import numpy as np
import jax.numpy as jnp


# Model function: use jnp
def model(x, a, b):
    return a * jnp.exp(-b * x)  # jnp


# Data generation: use np
x = np.linspace(0, 5, 100)  # np is fine for data
y = model(x, 2.5, 1.3)  # JAX auto-converts

# Fitting: works with both
popt, pcov = curve_fit(model, x, y, p0=[2, 1])

Issue: ConcretizationTypeError: Abstract tracer value

Cause: Using Python control flow in JIT-compiled function

Problem:

# Bad: Python if statement (not JIT-compatible)
def model(x, a, b, c):
    if a > 0:  # Error!
        return a * jnp.exp(-b * x)
    else:
        return c

Solution:

# Good: Use JAX control flow
def model(x, a, b, c):
    return jnp.where(a > 0, a * jnp.exp(-b * x), c)

Issue: TracerBoolConversionError

Cause: Using array values in boolean context

Problem:

# Bad
def model(x, a, b):
    if x[0] > 5:  # Error! Can't convert traced array to bool
        return a * x
    return b * x

Solution:

# Good
def model(x, a, b):
    return jnp.where(x > 5, a * x, b * x)

---

Diagnostic Flowchart

Fit fails?
  │
  ├─> ImportError/ModuleNotFoundError ──> Installation Issues
  │
  ├─> RuntimeError: "not found" ─────────> Convergence Problems
  │
  ├─> MemoryError/OOM ───────────────────> Memory Issues
  │
  ├─> Slow performance ──────────────────> Performance Issues
  │
  ├─> NaN/Inf in results ────────────────> Numerical Stability
  │
  ├─> TypeError (JAX/tracing) ───────────> JAX-Specific Issues
  │
  └─> Other ─────────────────────────────> API and Usage Errors

---

Getting Help

If this guide doesn’t resolve your issue:

1. Check documentation:

   • API Reference

   • Configuration Reference

   • Advanced Customization Guide

   • Performance Optimization Guide

2. Search GitHub issues:

   • https://github.com/imewei/NLSQ/issues

3. Create minimal reproducible example:

import numpy as np
import jax.numpy as jnp
from nlsq import curve_fit

# Minimal data
x = np.linspace(0, 5, 50)
y = 2.5 * np.exp(-1.3 * x) + 0.1 * np.random.randn(50)


# Minimal model
def model(x, a, b):
    return a * jnp.exp(-b * x)


# Your issue
popt, pcov = curve_fit(model, x, y, p0=[2, 1])  # Describe problem here

4. Report bug with:

   • Python version

   • JAX version (jax.__version__)

   • NLSQ version (nlsq.__version__)

   • GPU info (if applicable)

   • Minimal reproducible code

   • Full error traceback

---

Interactive Notebooks

Hands-on tutorials for debugging and troubleshooting:

• Troubleshooting Guide (25 min) - Debug convergence issues and common problems

• NLSQ Challenges (45 min) - Difficult optimization problems and solutions

• Hybrid Streaming API - Adaptive hybrid streaming overview

---

Quick Reference: Common Solutions

Problem	Quick Fix
Import error	pip install nlsq
No GPU	pip install "jax[cuda12_pip]"
Out of memory	Use fit_large_dataset()
Slow first fit	Use CurveFit() class
Convergence failure	Better p0, set bounds
NaN/Inf	Add numerical safeguards
Tracing error	Use jnp.where() not if
Wrong results	Check p0, visualize fit
Large covariance	Scale parameters
Slow with small data	Force CPU mode

---

Last Updated: 2025-10-07