Repository: scipy/scipy
File: scipy/_lib/_array_api_override.py (Python)

Code before:

       - `numpy.ma.MaskedArray`
       - `numpy.matrix`
       - NumPy arrays which do not have a boolean or numerical dtype
       - `scipy.sparse` arrays

    3. Coerce array-likes to NumPy arrays and check their dtype.
       Note that non-scalar array-likes can't be mixed with non-NumPy Array
       API objects; e.g.

       - `array_namespace([1, 2])` returns NumPy namespace;
       - `array_namespace(np.asarray([1, 2], [3, 4])` returns NumPy namespace;
       - `array_namespace(cp.asarray([1, 2], [3, 4])` raises an error.
    """
    if not SCIPY_ARRAY_API:
        # here we could wrap the namespace if needed
        return np_compat

    numpy_arrays = []
    api_arrays = []

    for array in arrays:
        arr_info = _validate_array_cls(type(array))
        if arr_info is _ArrayClsInfo.skip:
            pass
        elif arr_info is _ArrayClsInfo.is_numpy:
            dtype_info = _validate_numpy_dtype(array.dtype)
            if dtype_info is _DTypeInfo.maybe_jax:
                if is_jax_array(array):
                    api_arrays.append(array)  # JAX zero gradient array
                else:
                    # void dtype, but not a JAX zero gradient array
                    raise TypeError(f"An argument has dtype `{array.dtype!r}`; "
                        f"only boolean and numerical dtypes are supported.")
            else:
                numpy_arrays.append(array)
        elif arr_info is _ArrayClsInfo.no_info and is_array_api_obj(array):
            api_arrays.append(array)
        else:
            # list, tuple, or arbitrary object
            try:
                array = np.asanyarray(array)
            except TypeError:
                raise TypeError("An argument is neither array API compatible nor "
                                "coercible by NumPy.")
            _validate_numpy_dtype(array.dtype)
            numpy_arrays.append(array)

    # When there are exclusively NumPy and ArrayLikes, skip calling
    # array_api_compat.array_namespace for performance.
    if not api_arrays:

Review comment (suggestion):
We could drastically speed up the use case of ArrayLikes by pushing this validation to `_asarray`. However, that would require sanitizing all calls to `xp.asarray` throughout the codebase, so I'm leaving it to a follow-up.

Code after:
        return np_compat

    numpy_arrays = []
    api_arrays = []

    for array in arrays:
        arr_info = _validate_array_cls(type(array))
        if arr_info is _ArrayClsInfo.skip:
            pass

        elif arr_info is _ArrayClsInfo.numpy:
            if array.dtype.kind in 'iufcb':  # Numeric or bool
                numpy_arrays.append(array)
            elif array.dtype.kind == 'V' and is_jax_array(array):
                # Special case for JAX zero gradient arrays;
                # see array_api_compat._common._helpers._is_jax_zero_gradient_array
                api_arrays.append(array)  # JAX zero gradient array
            else:
                raise TypeError(f"An argument has dtype `{array.dtype!r}`; "
                                "only boolean and numerical dtypes are supported.")

        elif arr_info is _ArrayClsInfo.unknown and is_array_api_obj(array):
            api_arrays.append(array)

        else:
            # list, tuple, or arbitrary object
            try:
                array = np.asanyarray(array)
            except TypeError:
                raise TypeError("An argument is neither array API compatible nor "
                                "coercible by NumPy.")
            if array.dtype.kind not in 'iufcb':  # Numeric or bool
                raise TypeError(f"An argument has dtype `{array.dtype!r}`; "
                                "only boolean and numerical dtypes are supported.")
            numpy_arrays.append(array)

    # When there are exclusively NumPy and ArrayLikes, skip calling
    # array_api_compat.array_namespace for performance.
    if not api_arrays:
        return np_compat

    # In case of mix of NumPy/ArrayLike and non-NumPy Array API arrays,
    # let array_api_compat.array_namespace raise an error.
    return array_api_compat.array_namespace(*numpy_arrays, *api_arrays)