Skip to content

Improve how the SciPy BLAS library is found #633

Open
@dime10

Description

@dime10

Catalyst support for certain JAX functions relies on a BLAS runtime library being available, specifically the OpenBLAS library shipped with SciPy. The relevant source code is located in the Python compiler driver:

if platform.system() == "Linux":
file_path_within_package = "../scipy.libs/"
file_extension = ".so"
elif platform.system() == "Darwin": # pragma: nocover
file_path_within_package = ".dylibs/"
file_extension = ".dylib"
package_name = "scipy"
scipy_package = importlib.util.find_spec(package_name)
package_directory = path.dirname(scipy_package.origin)
scipy_lib_path = path.join(package_directory, file_path_within_package)
file_prefix = "libopenblas"
search_pattern = path.join(scipy_lib_path, f"{file_prefix}*{file_extension}")
openblas_so_file = glob.glob(search_pattern)[0]
openblas_lib_name = path.basename(openblas_so_file)[3 : -len(file_extension)]

This code snippet is brittle and susceptible to different packaging configurations (for instance the assumption of the folder locations relative to the scipy package location.

At minimum we should raise a user friendly error message that there is a problem with the scipy installation.

Ideally, we would also find ways to make this more robust on the various systems, and also consider conda installations.

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions