Add functionality for the data_dependent_shapes flag

Author: asmeurer

array_api_strict/_array_object.py

@@ -32,6 +32,7 @@
     _result_type,
     _dtype_categories,
 )
+from ._flags import get_array_api_strict_flags
 
 from typing import TYPE_CHECKING, Optional, Tuple, Union, Any, SupportsIndex
 import types
@@ -427,13 +428,17 @@ def _validate_index(self, key):
                                 "the Array API)"
                             )
             elif isinstance(i, Array):
-                if i.dtype in _boolean_dtypes and len(_key) != 1:
-                    assert isinstance(key, tuple)  # sanity check
-                    raise IndexError(
-                        f"Single-axes index {i} is a boolean array and "
-                        f"{len(key)=}, but masking is only specified in the "
-                        "Array API when the array is the sole index."
-                    )
+                if i.dtype in _boolean_dtypes:
+                    if len(_key) != 1:
+                        assert isinstance(key, tuple)  # sanity check
+                        raise IndexError(
+                            f"Single-axes index {i} is a boolean array and "
+                            f"{len(key)=}, but masking is only specified in the "
+                            "Array API when the array is the sole index."
+                        )
+                    if not get_array_api_strict_flags()['data_dependent_shapes']:
+                        raise RuntimeError("Boolean array indexing (masking) requires data-dependent shapes, but the data_dependent_shapes flag has been disabled for array-api-strict")
+
                 elif i.dtype in _integer_dtypes and i.ndim != 0:
                     raise IndexError(
                         f"Single-axes index {i} is a non-zero-dimensional "

array_api_strict/_flags.py

@@ -265,3 +265,11 @@ def set_flags_from_environment():
         )
 
 set_flags_from_environment()
+
+def requires_data_dependent_shapes(func):
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        if not DATA_DEPENDENT_SHAPES:
+            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

array_api_strict/_searching_functions.py

@@ -2,6 +2,7 @@
 
 from ._array_object import Array
 from ._dtypes import _result_type, _real_numeric_dtypes
+from ._flags import requires_data_dependent_shapes
 
 from typing import Optional, Tuple
 
@@ -30,6 +31,7 @@ def argmin(x: Array, /, *, axis: Optional[int] = None, keepdims: bool = False) -
     return Array._new(np.asarray(np.argmin(x._array, axis=axis, keepdims=keepdims)))
 
 
+@requires_data_dependent_shapes
 def nonzero(x: Array, /) -> Tuple[Array, ...]:
     """
     Array API compatible wrapper for :py:func:`np.nonzero <numpy.nonzero>`.

array_api_strict/_set_functions.py

@@ -2,6 +2,8 @@
 
 from ._array_object import Array
 
+from ._flags import requires_data_dependent_shapes
+
 from typing import NamedTuple
 
 import numpy as np
@@ -35,6 +37,7 @@ class UniqueInverseResult(NamedTuple):
     inverse_indices: Array
 
 
+@requires_data_dependent_shapes
 def unique_all(x: Array, /) -> UniqueAllResult:
     """
     Array API compatible wrapper for :py:func:`np.unique <numpy.unique>`.
@@ -59,6 +62,7 @@ def unique_all(x: Array, /) -> UniqueAllResult:
     )
 
 
+@requires_data_dependent_shapes
 def unique_counts(x: Array, /) -> UniqueCountsResult:
     res = np.unique(
         x._array,
@@ -71,6 +75,7 @@ def unique_counts(x: Array, /) -> UniqueCountsResult:
     return UniqueCountsResult(*[Array._new(i) for i in res])
 
 
+@requires_data_dependent_shapes
 def unique_inverse(x: Array, /) -> UniqueInverseResult:
     """
     Array API compatible wrapper for :py:func:`np.unique <numpy.unique>`.
@@ -90,6 +95,7 @@ def unique_inverse(x: Array, /) -> UniqueInverseResult:
     return UniqueInverseResult(Array._new(values), Array._new(inverse_indices))
 
 
+@requires_data_dependent_shapes
 def unique_values(x: Array, /) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.unique <numpy.unique>`.