array-api-strict 2.2__py3-none-any.whl → 2.3.1__py3-none-any.whl

array_api_strict/__init__.py

@@ -293,9 +293,9 @@ from ._manipulation_functions import (
 
 __all__ += ["concat", "expand_dims", "flip", "moveaxis", "permute_dims", "repeat", "reshape", "roll", "squeeze", "stack", "tile", "unstack"]
 
-from ._searching_functions import argmax, argmin, nonzero, searchsorted, where
+from ._searching_functions import argmax, argmin, nonzero, count_nonzero, searchsorted, where
 
-__all__ += ["argmax", "argmin", "nonzero", "searchsorted", "where"]
+__all__ += ["argmax", "argmin", "nonzero", "count_nonzero", "searchsorted", "where"]
 
 from ._set_functions import unique_all, unique_counts, unique_inverse, unique_values
 
@@ -305,9 +305,9 @@ from ._sorting_functions import argsort, sort
 
 __all__ += ["argsort", "sort"]
 
-from ._statistical_functions import cumulative_sum, max, mean, min, prod, std, sum, var
+from ._statistical_functions import cumulative_sum, cumulative_prod, max, mean, min, prod, std, sum, var
 
-__all__ += ["cumulative_sum", "max", "mean", "min", "prod", "std", "sum", "var"]
+__all__ += ["cumulative_sum", "cumulative_prod", "max", "mean", "min", "prod", "std", "sum", "var"]
 
 from ._utility_functions import all, any, diff
 
@@ -327,9 +327,12 @@ from ._flags import (
 
 __all__ += ['set_array_api_strict_flags', 'get_array_api_strict_flags', 'reset_array_api_strict_flags', 'ArrayAPIStrictFlags']
 
-from . import _version
-__version__ = _version.get_versions()['version']
-del _version
+try:
+    from . import _version
+    __version__ = _version.__version__
+    del _version
+except ImportError:
+    __version__ = "unknown"
 
 
 # Extensions can be enabled or disabled dynamically. In order to make

array_api_strict/_array_object.py

@@ -26,10 +26,12 @@ from ._dtypes import (
     _integer_dtypes,
     _integer_or_boolean_dtypes,
     _floating_dtypes,
+    _real_floating_dtypes,
     _complex_floating_dtypes,
     _numeric_dtypes,
     _result_type,
     _dtype_categories,
+    _real_to_complex_map,
 )
 from ._flags import get_array_api_strict_flags, set_array_api_strict_flags
 
@@ -126,12 +128,6 @@ class Array:
     # These functions are not required by the spec, but are implemented for
     # the sake of usability.
 
-    def __str__(self: Array, /) -> str:
-        """
-        Performs the operation __str__.
-        """
-        return self._array.__str__().replace("array", "Array")
-
     def __repr__(self: Array, /) -> str:
         """
         Performs the operation __repr__.
@@ -149,6 +145,8 @@ class Array:
             mid = np.array2string(self._array, separator=', ', prefix=prefix, suffix=suffix)
         return prefix + mid + suffix
 
+    __str__ = __repr__
+
     # In the future, _allow_array will be set to False, which will disallow
     # __array__. This means calling `np.func()` on an array_api_strict array
     # will give an error. If we don't explicitly disallow it, NumPy defaults
@@ -247,6 +245,7 @@ class Array:
         """
         from ._data_type_functions import iinfo
 
+        target_dtype = self.dtype
         # Note: Only Python scalar types that match the array dtype are
         # allowed.
         if isinstance(scalar, bool):
@@ -272,10 +271,13 @@ class Array:
                     "Python float scalars can only be promoted with floating-point arrays."
                 )
         elif isinstance(scalar, complex):
-            if self.dtype not in _complex_floating_dtypes:
+            if self.dtype not in _floating_dtypes:
                 raise TypeError(
-                    "Python complex scalars can only be promoted with complex floating-point arrays."
+                    "Python complex scalars can only be promoted with floating-point arrays."
                 )
+            # 1j * array(floating) is allowed
+            if self.dtype in _real_floating_dtypes:
+                target_dtype = _real_to_complex_map[self.dtype]
         else:
             raise TypeError("'scalar' must be a Python scalar")
 
@@ -286,7 +288,7 @@ class Array:
         # behavior for integers within the bounds of the integer dtype.
         # Outside of those bounds we use the default NumPy behavior (either
         # cast or raise OverflowError).
-        return Array._new(np.array(scalar, dtype=self.dtype._np_dtype), device=self.device)
+        return Array._new(np.array(scalar, dtype=target_dtype._np_dtype), device=self.device)
 
     @staticmethod
     def _normalize_two_args(x1, x2) -> Tuple[Array, Array]:
@@ -325,7 +327,7 @@ class Array:
 
     # Note: A large fraction of allowed indices are disallowed here (see the
     # docstring below)
-    def _validate_index(self, key):
+    def _validate_index(self, key, op="getitem"):
         """
         Validate an index according to the array API.
 
@@ -388,11 +390,16 @@ class Array:
                     "zero-dimensional integer arrays and boolean arrays "
                     "are specified in the Array API."
                 )
+            if op == "setitem":
+                if isinstance(i, Array) and i.dtype in _integer_dtypes:
+                    raise IndexError("Fancy indexing __setitem__ is not supported.")
 
         nonexpanding_key = []
         single_axes = []
         n_ellipsis = 0
         key_has_mask = False
+        key_has_index_array = False
+        key_has_slices = False
         for i in _key:
             if i is not None:
                 nonexpanding_key.append(i)
@@ -401,6 +408,8 @@ class Array:
                 if isinstance(i, Array):
                     if i.dtype in _boolean_dtypes:
                         key_has_mask = True
+                    elif i.dtype in _integer_dtypes:
+                        key_has_index_array = True
                     single_axes.append(i)
                 else:
                     # i must not be an array here, to avoid elementwise equals
@@ -408,6 +417,8 @@ class Array:
                         n_ellipsis += 1
                     else:
                         single_axes.append(i)
+                        if isinstance(i, slice):
+                            key_has_slices = True
 
         n_single_axes = len(single_axes)
         if n_ellipsis > 1:
@@ -425,6 +436,12 @@ class Array:
                     "specified in the Array API."
                 )
 
+        if (key_has_index_array and (n_ellipsis > 0 or key_has_slices or key_has_mask)):
+            raise IndexError(
+                "Integer index arrays are only allowed with integer indices; "
+                f"got {key}."
+            )
+
         if n_ellipsis == 0:
             indexed_shape = self.shape
         else:
@@ -481,14 +498,9 @@ class Array:
                             "Array API when the array is the sole index."
                         )
                     if not get_array_api_strict_flags()['boolean_indexing']:
-                        raise RuntimeError("The boolean_indexing 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 "
-                        "integer array, but advanced integer indexing is not "
-                        "specified in the Array API."
-                    )
+                        raise RuntimeError(
+                            "The boolean_indexing flag has been disabled for array-api-strict"
+                        )
             elif isinstance(i, tuple):
                 raise IndexError(
                     f"Single-axes index {i} is a tuple, but nested tuple "
@@ -900,7 +912,7 @@ class Array:
         """
         # Note: Only indices required by the spec are allowed. See the
         # docstring of _validate_index
-        self._validate_index(key)
+        self._validate_index(key, op="setitem")
         if isinstance(key, Array):
             # Indexing self._array with array_api_strict arrays can be erroneous
             key = key._array

array_api_strict/_creation_functions.py

@@ -226,6 +226,8 @@ def from_dlpack(
     # Going to wait for upstream numpy support
     if device is not _default:
         _check_device(device)
+    else:
+        device = None
     if copy not in [_default, None]:
         raise NotImplementedError("The copy argument to from_dlpack is not yet implemented")
 

array_api_strict/_data_type_functions.py

@@ -42,6 +42,13 @@ def astype(
 
     if not copy and dtype == x.dtype:
         return x
+
+    if isdtype(x.dtype, 'complex floating') and not isdtype(dtype, 'complex floating'):
+        raise TypeError(
+            f'The Array API standard stipulates that casting {x.dtype} to {dtype} should not be permitted. '
+             'array-api-strict thus prohibits this conversion.'
+        )
+
     return Array._new(x._array.astype(dtype=dtype._np_dtype, copy=copy), device=device)
 
 
@@ -160,6 +167,9 @@ def isdtype(
     https://data-apis.org/array-api/latest/API_specification/generated/array_api.isdtype.html
     for more details
     """
+    if not isinstance(dtype, _DType):
+       raise TypeError(f"'dtype' must be a dtype, not a {type(dtype)!r}")
+
     if isinstance(kind, tuple):
         # Disallow nested tuples
         if any(isinstance(k, tuple) for k in kind):
@@ -187,7 +197,7 @@ def isdtype(
     else:
         raise TypeError(f"'kind' must be a dtype, str, or tuple of dtypes and strs, not {type(kind).__name__}")
 
-def result_type(*arrays_and_dtypes: Union[Array, Dtype]) -> Dtype:
+def result_type(*arrays_and_dtypes: Union[Array, Dtype, int, float, complex, bool]) -> Dtype:
     """
     Array API compatible wrapper for :py:func:`np.result_type <numpy.result_type>`.
 
@@ -198,19 +208,40 @@ def result_type(*arrays_and_dtypes: Union[Array, Dtype]) -> Dtype:
     # too many extra type promotions like int64 + uint64 -> float64, and does
     # value-based casting on scalar arrays.
     A = []
+    scalars = []
     for a in arrays_and_dtypes:
         if isinstance(a, Array):
             a = a.dtype
+        elif isinstance(a, (bool, int, float, complex)):
+            scalars.append(a)
         elif isinstance(a, np.ndarray) or a not in _all_dtypes:
             raise TypeError("result_type() inputs must be array_api arrays or dtypes")
         A.append(a)
 
+    # remove python scalars
+    A = [a for a in A if not isinstance(a, (bool, int, float, complex))]
+
     if len(A) == 0:
         raise ValueError("at least one array or dtype is required")
     elif len(A) == 1:
-        return A[0]
+        result = A[0]
     else:
         t = A[0]
         for t2 in A[1:]:
             t = _result_type(t, t2)
-        return t
+        result = t
+
+    if len(scalars) == 0:
+        return result
+
+    if get_array_api_strict_flags()['api_version'] <= '2023.12':
+        raise TypeError("result_type() inputs must be array_api arrays or dtypes")
+
+    # promote python scalars given the result_type for all arrays/dtypes
+    from ._creation_functions import empty
+    arr = empty(1, dtype=result)
+    for s in scalars:
+        x = arr._promote_scalar(s)
+        result = _result_type(x.dtype, result)
+
+    return result

array_api_strict/_dtypes.py

@@ -126,6 +126,7 @@ _dtype_categories = {
     "floating-point": _floating_dtypes,
 }
 
+_real_to_complex_map = {float32: complex64, float64: complex128}
 
 # Note: the spec defines a restricted type promotion table compared to NumPy.
 # In particular, cross-kind promotions like integer + float or boolean +