RFC: add Array API Extension support function

[purepani]

Right now, the dask library has a linalg attribute, but doesn't currently implement the entire linalg extension spec. The lack of full support is mentioned here: https://data-apis.org/array-api-compat/supported-array-libraries.html

This means that an array api library consumer can't do

if hasattr(xp, 'linalg'):
    xp.linalg.cross(x, y)

as recommended here,
and work with dask. I suppose this is hard to change for existing array libraries that have such an attribute already, at least without creating a whole new namespace.
Because of this, I propose adding a supports_extension function and/or supported_extensions attribute to the spec, so that

from array_api_compat import array_namespace

def cross(x, y):
    xp = array_namespace(x, y)
    if xp.supports_extension('linalg'):
        return xp.linalg.cross(x, y)

is possible.

Alternatively:

from array_api_compat import array_namespace

def cross(x, y):
    xp = array_namespace(x, y)
    if 'linalg' in xp.extensions:
        return xp.linalg.cross(x, y)

In the meantime, I created data-apis/array-api-compat#332 for array-api-compat to deal with this problem, but I think this makes more sense as part of the spec.

[rgommers]

Thanks for the suggestion @purepani. The problem is clear, so a few thoughts on potential solutions:

• A top-level public function like supports_extension for this kind of introspection won't fly. That requires all array libraries to add something to their public namespace that isn't really used.
• Instead, we have a better hidden introspection method: https://data-apis.org/array-api/2024.12/API_specification/inspection.html. Adding supported extensions could be added there, if we want a function exactly like you propose.
• That said, this is a fairly straightforward case of "spec is only partially implemented". There are a couple of established ways of dealing with that:
  • Adding the missing function (cross in this case) to array-api-compat
  • Dealing with it locally in your library with if is_dask(xp): or if hasattr(xp.linalg, 'cross'): rather than if xp.supports_extension('linalg'):` (would be more concise and lower-overhead).
  • And of course: implementing the cross function in dask to address it comprehensively

I think that the last three three bullets can all be done, and then there is no need for this introspection. It's also a dangerous kind of inspection because (a) it will break code that already does work (most of the functionality is present after all: https://docs.dask.org/en/stable/array-api.html#linear-algebra), and (b) a new standard version could add more functions to linalg and then what does a library do if they supported all the old ones but are blocked on adding the new one for some reason?

All the existing introspection with __array_namespace_info__ is straightforward and unlikely to change for new library or standard version. I don't think it's appropriate to start adding it for more volatile information that can already be accessed with hasattr.

[purepani]

I've been using if hasattr(xp.linalg, 'cross'), but to me that just seems at least extremely cumbersome to do it for every function, especially for a linalg heavy library, though the one I'm working on isn't super heavy, one that is can imagine having to check for these everywhere is somewhat frustrating and adds a lot of noise; additionally, having to check for each individual linalg function you're going to use seems counter to the point of having a standard.

I'll go ahead and make an issue in dask as well about implementing these as well.
I guess the question is, should dask be considered to be not compliant with this part of the array API spec since it doesn't fully implement linalg, but doesn't stop you from checking for the submodule. If the answer is that it isn't compliant and it needs to be fixed, then I don't think there's an issue, but if the answer is that we still want to call it compliant as it currently is, then I think having something in the inspection module makes sense.

To answer the thought about the linalg module being stuck at a version in the event something is added to the inspection module, maybe it makes sense to version the extensions separately from the full library, so that 2025.12 could have a 2024.12 linalg extension. This opens a whole new can of worms, so maybe it's not a super good idea, but it is an idea.

As a note, I wasn't super clear. This isn't just cross, but something more like at least 10 or 15 functions, at least if https://data-apis.org/array-api-compat/supported-array-libraries.html#dask is correct, so this isn't going to be a super quick fix to implement everything, though it could be done(though my naive understanding of the library is that some of these are difficult to implement cleanly into dask). To me it makes more sense to just say that the entire linalg module isn't implemented, and then special case for dask whenever a function is implemented, rather than special casing for every function that isn't currently implemented.

[rgommers]

If the answer is that it isn't compliant and it needs to be fixed, then I don't think there's an issue

Yes, that's the answer here.

but to me that just seems at least extremely cumbersome to do it for every function

Only for the ones you know are missing, right? And if the function that is missing cannot be easily supported through another code path, you will anyway have to do:

def func_relying_on_cross(x, y):
    xp = array_namespace(x, y)
    if is_dask(xp):
        raise TypeError("`func_relying_on_cross` does not support Dask arrays, because Dask misses `linalg.cross`")        

This is how we approach it in SciPy at least. I don't really see a better way.

having to check for each individual linalg function you're going to use seems counter to the point of having a standard.

Agreed. This is not really different though than missing functions in the main namespace. It has always been, and still is, a challenge to deal with libraries that only partially support the standard. There really is no painless way around that, and that is also true for how other standards work (e.g., because MSVC doesn't support C99 complex types, no C library one can use those types without a cumbersome workaround if they care about MSVC support).

If a library has too many holes in it, it probably shouldn't declare support. Or the array-consuming library should exclude it from its list of supported libraries. In your example: if half of linalg is missing in Dask, I would not expect any linalg-heavy library that needs that missing functionality to support Dask until most of the gaps are filled in.

For Dask today, I think enough functions are supported that it declaring support is fine, but it does need workarounds sometimes.