nlsq.device.check_gpu_availability

nlsq.device.check_gpu_availability(warn=True)[source]

Check if GPU is available but not being used by JAX.

Prints a helpful warning if GPU hardware and system CUDA are detected but JAX is running in CPU-only mode.

Parameters:

warn (bool, default=True) – If True, print warning when GPU available but not used.

Returns:

  • bool – True if GPU is being used by JAX, False otherwise.

  • Environment Variables

  • ———————

  • NLSQ_SKIP_GPU_CHECK (str, optional) – Set to “1”, “true”, or “yes” (case-insensitive) to suppress GPU warnings. Useful for CI/CD pipelines or users who intentionally use CPU-only JAX.

Return type:

bool

Notes

This check runs automatically on import but has minimal overhead: - Subprocess call to nvidia-smi (~5ms) - JAX device query (~1ms) - Only prints warning when mismatch detected - Silent failures prevent disruption