NVIDIA cuSPARSE (CUDA 12)

import os
try:
    import cupy as cp
    import cupy.sparse as cps
    import numpy as np

    # Check for CUDA device
    if cp.cuda.is_available():
        print(f"CUDA is available. CuPy version: {cp.__version__}")
        print(f"CUDA Device Name: {cp.cuda.Device().name}")

        # Create a sparse matrix on CPU (SciPy format)
        row = np.array([0, 1, 2, 0])
        col = np.array([0, 1, 2, 1])
        data = np.array([1, 2, 3, 4])
        shape = (3, 3)
        sparse_cpu = cps.csr_matrix((data, (row, col)), shape=shape)
        print("\nCPU Sparse Matrix:\n", sparse_cpu.toarray())

        # Transfer to GPU and perform a sparse matrix-vector multiplication
        sparse_gpu = cps.csr_matrix(sparse_cpu, dtype=cp.float32) # cuSPARSE generally works with float32/64
        vector_gpu = cp.array([1.0, 2.0, 3.0], dtype=cp.float32)

        result_gpu = sparse_gpu @ vector_gpu
        print("\nGPU Sparse Matrix-Vector Product (using cuSPARSE via CuPy):\n", result_gpu)

    else:
        print("CUDA is not available. Please ensure a compatible NVIDIA GPU and driver are installed.")
        print("You may need to install cupy-cuda12x manually if using a specific CUDA version.")

except ImportError:
    print("CuPy is not installed. To run this example, install CuPy compatible with CUDA 12:")
    print("pip install cupy-cuda12x")
except Exception as e:
    print(f"An error occurred: {e}")