Respect the extension flag in linalg and fft

This behavior still needs to be tested.

This required moving the linalg functions that are also in the main namespace
so that they can still work there even when the linalg extension is disabled.

The way I've decided to implement this is that the functions will not raise an
exception until they are called. It would probably be more convenient for
users if they raised an attribute error, or if the extension namespace itself
did, like it would in a real library without the given extension. But the
implementation for this would be a lot more complicated and didn't really feel
worth it to me.

Author: asmeurer

array_api_strict/__init__.py

@@ -244,7 +244,7 @@
 
 __all__ += ["linalg"]
 
-from .linalg import matmul, tensordot, matrix_transpose, vecdot
+from ._linear_algebra_functions import matmul, tensordot, matrix_transpose, vecdot
 
 __all__ += ["matmul", "tensordot", "matrix_transpose", "vecdot"]
 

array_api_strict/_flags.py

@@ -273,3 +273,17 @@ def wrapper(*args, **kwargs):
             raise RuntimeError(f"The function {func.__name__} requires data-dependent shapes, but the data_dependent_shapes flag has been disabled for array-api-strict")
         return func(*args, **kwargs)
     return wrapper
+
+def requires_extension(extension):
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            if extension not in ENABLED_EXTENSIONS:
+                if extension == 'linalg' \
+                   and func.__name__ in ['matmul', 'tensordot',
+                                         'matrix_transpose', 'vecdot']:
+                    raise RuntimeError(f"The linalg extension has been disabled for array-api-strict. However, {func.__name__} is also present in the main array_api_strict namespace and may be used from there.")
+                raise RuntimeError(f"The function {func.__name__} requires the {extension} extension, but it has been disabled for array-api-strict")
+            return func(*args, **kwargs)
+        return wrapper
+    return decorator

array_api_strict/_linear_algebra_functions.py

@@ -0,0 +1,68 @@
+"""
+These functions are all also defined in the linalg extension, but we include
+them here with wrappers in linalg so that the wrappers can be disabled if the
+linalg extension is disabled in the flags.
+
+"""
+
+from __future__ import annotations
+
+from ._dtypes import _numeric_dtypes
+
+from ._array_object import Array
+
+from typing import TYPE_CHECKING
+if TYPE_CHECKING:
+    from ._typing import Sequence, Tuple, Union
+
+import numpy.linalg
+import numpy as np
+
+# Note: matmul is the numpy top-level namespace but not in np.linalg
+def matmul(x1: Array, x2: Array, /) -> Array:
+    """
+    Array API compatible wrapper for :py:func:`np.matmul <numpy.matmul>`.
+
+    See its docstring for more information.
+    """
+    # Note: the restriction to numeric dtypes only is different from
+    # np.matmul.
+    if x1.dtype not in _numeric_dtypes or x2.dtype not in _numeric_dtypes:
+        raise TypeError('Only numeric dtypes are allowed in matmul')
+
+    return Array._new(np.matmul(x1._array, x2._array))
+
+# Note: tensordot is the numpy top-level namespace but not in np.linalg
+
+# Note: axes must be a tuple, unlike np.tensordot where it can be an array or array-like.
+def tensordot(x1: Array, x2: Array, /, *, axes: Union[int, Tuple[Sequence[int], Sequence[int]]] = 2) -> Array:
+    # Note: the restriction to numeric dtypes only is different from
+    # np.tensordot.
+    if x1.dtype not in _numeric_dtypes or x2.dtype not in _numeric_dtypes:
+        raise TypeError('Only numeric dtypes are allowed in tensordot')
+
+    return Array._new(np.tensordot(x1._array, x2._array, axes=axes))
+
+# Note: this function is new in the array API spec. Unlike transpose, it only
+# transposes the last two axes.
+def matrix_transpose(x: Array, /) -> Array:
+    if x.ndim < 2:
+        raise ValueError("x must be at least 2-dimensional for matrix_transpose")
+    return Array._new(np.swapaxes(x._array, -1, -2))
+
+# Note: vecdot is not in NumPy
+def vecdot(x1: Array, x2: Array, /, *, axis: int = -1) -> Array:
+    if x1.dtype not in _numeric_dtypes or x2.dtype not in _numeric_dtypes:
+        raise TypeError('Only numeric dtypes are allowed in vecdot')
+    ndim = max(x1.ndim, x2.ndim)
+    x1_shape = (1,)*(ndim - x1.ndim) + tuple(x1.shape)
+    x2_shape = (1,)*(ndim - x2.ndim) + tuple(x2.shape)
+    if x1_shape[axis] != x2_shape[axis]:
+        raise ValueError("x1 and x2 must have the same size along the given axis")
+
+    x1_, x2_ = np.broadcast_arrays(x1._array, x2._array)
+    x1_ = np.moveaxis(x1_, axis, -1)
+    x2_ = np.moveaxis(x2_, axis, -1)
+
+    res = x1_[..., None, :] @ x2_[..., None]
+    return Array._new(res[..., 0, 0])

array_api_strict/fft.py

@@ -15,9 +15,11 @@
 )
 from ._array_object import Array, CPU_DEVICE
 from ._data_type_functions import astype
+from ._flags import requires_extension
 
 import numpy as np
 
+@requires_extension('fft')
 def fft(
     x: Array,
     /,
@@ -40,6 +42,7 @@ def fft(
         return astype(res, complex64)
     return res
 
+@requires_extension('fft')
 def ifft(
     x: Array,
     /,
@@ -62,6 +65,7 @@ def ifft(
         return astype(res, complex64)
     return res
 
+@requires_extension('fft')
 def fftn(
     x: Array,
     /,
@@ -84,6 +88,7 @@ def fftn(
         return astype(res, complex64)
     return res
 
+@requires_extension('fft')
 def ifftn(
     x: Array,
     /,
@@ -106,6 +111,7 @@ def ifftn(
         return astype(res, complex64)
     return res
 
+@requires_extension('fft')
 def rfft(
     x: Array,
     /,
@@ -128,6 +134,7 @@ def rfft(
         return astype(res, complex64)
     return res
 
+@requires_extension('fft')
 def irfft(
     x: Array,
     /,
@@ -150,6 +157,7 @@ def irfft(
         return astype(res, float32)
     return res
 
+@requires_extension('fft')
 def rfftn(
     x: Array,
     /,
@@ -172,6 +180,7 @@ def rfftn(
         return astype(res, complex64)
     return res
 
+@requires_extension('fft')
 def irfftn(
     x: Array,
     /,
@@ -194,6 +203,7 @@ def irfftn(
         return astype(res, float32)
     return res
 
+@requires_extension('fft')
 def hfft(
     x: Array,
     /,
@@ -216,6 +226,7 @@ def hfft(
         return astype(res, float32)
     return res
 
+@requires_extension('fft')
 def ihfft(
     x: Array,
     /,
@@ -238,6 +249,7 @@ def ihfft(
         return astype(res, complex64)
     return res
 
+@requires_extension('fft')
 def fftfreq(n: int, /, *, d: float = 1.0, device: Optional[Device] = None) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.fft.fftfreq <numpy.fft.fftfreq>`.
@@ -248,6 +260,7 @@ def fftfreq(n: int, /, *, d: float = 1.0, device: Optional[Device] = None) -> Ar
         raise ValueError(f"Unsupported device {device!r}")
     return Array._new(np.fft.fftfreq(n, d=d))
 
+@requires_extension('fft')
 def rfftfreq(n: int, /, *, d: float = 1.0, device: Optional[Device] = None) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.fft.rfftfreq <numpy.fft.rfftfreq>`.
@@ -258,6 +271,7 @@ def rfftfreq(n: int, /, *, d: float = 1.0, device: Optional[Device] = None) -> A
         raise ValueError(f"Unsupported device {device!r}")
     return Array._new(np.fft.rfftfreq(n, d=d))
 
+@requires_extension('fft')
 def fftshift(x: Array, /, *, axes: Union[int, Sequence[int]] = None) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.fft.fftshift <numpy.fft.fftshift>`.
@@ -268,6 +282,7 @@ def fftshift(x: Array, /, *, axes: Union[int, Sequence[int]] = None) -> Array:
         raise TypeError("Only floating-point dtypes are allowed in fftshift")
     return Array._new(np.fft.fftshift(x._array, axes=axes))
 
+@requires_extension('fft')
 def ifftshift(x: Array, /, *, axes: Union[int, Sequence[int]] = None) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.fft.ifftshift <numpy.fft.ifftshift>`.