Rename numpy._array_api to numpy.array_api

Instead of the leading underscore, the experimentalness of the module will be
indicated by omitting a warning on import. That we, we do not have to change
the API from underscore to no underscore when the module is no longer
experimental.

1 files changed, 250 insertions, 0 deletions

diff --git a/numpy/array_api/tests/test_array_object.py b/numpy/array_api/tests/test_array_object.py
new file mode 100644
index 000000000..22078bbee
--- /dev/null
+++ b/numpy/array_api/tests/test_array_object.py
@@ -0,0 +1,250 @@
+from numpy.testing import assert_raises
+import numpy as np
+
+from .. import ones, asarray, result_type
+from .._dtypes import (_all_dtypes, _boolean_dtypes, _floating_dtypes,
+                       _integer_dtypes, _integer_or_boolean_dtypes,
+                       _numeric_dtypes, int8, int16, int32, int64, uint64)
+
+def test_validate_index():
+    # The indexing tests in the official array API test suite test that the
+    # array object correctly handles the subset of indices that are required
+    # by the spec. But the NumPy array API implementation specifically
+    # disallows any index not required by the spec, via Array._validate_index.
+    # This test focuses on testing that non-valid indices are correctly
+    # rejected. See
+    # https://data-apis.org/array-api/latest/API_specification/indexing.html
+    # and the docstring of Array._validate_index for the exact indexing
+    # behavior that should be allowed. This does not test indices that are
+    # already invalid in NumPy itself because Array will generally just pass
+    # such indices directly to the underlying np.ndarray.
+
+    a = ones((3, 4))
+
+    # Out of bounds slices are not allowed
+    assert_raises(IndexError, lambda: a[:4])
+    assert_raises(IndexError, lambda: a[:-4])
+    assert_raises(IndexError, lambda: a[:3:-1])
+    assert_raises(IndexError, lambda: a[:-5:-1])
+    assert_raises(IndexError, lambda: a[3:])
+    assert_raises(IndexError, lambda: a[-4:])
+    assert_raises(IndexError, lambda: a[3::-1])
+    assert_raises(IndexError, lambda: a[-4::-1])
+
+    assert_raises(IndexError, lambda: a[...,:5])
+    assert_raises(IndexError, lambda: a[...,:-5])
+    assert_raises(IndexError, lambda: a[...,:4:-1])
+    assert_raises(IndexError, lambda: a[...,:-6:-1])
+    assert_raises(IndexError, lambda: a[...,4:])
+    assert_raises(IndexError, lambda: a[...,-5:])
+    assert_raises(IndexError, lambda: a[...,4::-1])
+    assert_raises(IndexError, lambda: a[...,-5::-1])
+
+    # Boolean indices cannot be part of a larger tuple index
+    assert_raises(IndexError, lambda: a[a[:,0]==1,0])
+    assert_raises(IndexError, lambda: a[a[:,0]==1,...])
+    assert_raises(IndexError, lambda: a[..., a[0]==1])
+    assert_raises(IndexError, lambda: a[[True, True, True]])
+    assert_raises(IndexError, lambda: a[(True, True, True),])
+
+    # Integer array indices are not allowed (except for 0-D)
+    idx = asarray([[0, 1]])
+    assert_raises(IndexError, lambda: a[idx])
+    assert_raises(IndexError, lambda: a[idx,])
+    assert_raises(IndexError, lambda: a[[0, 1]])
+    assert_raises(IndexError, lambda: a[(0, 1), (0, 1)])
+    assert_raises(IndexError, lambda: a[[0, 1]])
+    assert_raises(IndexError, lambda: a[np.array([[0, 1]])])
+
+    # np.newaxis is not allowed
+    assert_raises(IndexError, lambda: a[None])
+    assert_raises(IndexError, lambda: a[None, ...])
+    assert_raises(IndexError, lambda: a[..., None])
+
+def test_operators():
+    # For every operator, we test that it works for the required type
+    # combinations and raises TypeError otherwise
+    binary_op_dtypes ={
+        '__add__': 'numeric',
+        '__and__': 'integer_or_boolean',
+        '__eq__': 'all',
+        '__floordiv__': 'numeric',
+        '__ge__': 'numeric',
+        '__gt__': 'numeric',
+        '__le__': 'numeric',
+        '__lshift__': 'integer',
+        '__lt__': 'numeric',
+        '__mod__': 'numeric',
+        '__mul__': 'numeric',
+        '__ne__': 'all',
+        '__or__': 'integer_or_boolean',
+        '__pow__': 'floating',
+        '__rshift__': 'integer',
+        '__sub__': 'numeric',
+        '__truediv__': 'floating',
+        '__xor__': 'integer_or_boolean',
+    }
+
+    # Recompute each time because of in-place ops
+    def _array_vals():
+        for d in _integer_dtypes:
+            yield asarray(1, dtype=d)
+        for d in _boolean_dtypes:
+            yield asarray(False, dtype=d)
+        for d in _floating_dtypes:
+            yield asarray(1., dtype=d)
+
+    for op, dtypes in binary_op_dtypes.items():
+        ops = [op]
+        if op not in ['__eq__', '__ne__', '__le__', '__ge__', '__lt__', '__gt__']:
+            rop = '__r' + op[2:]
+            iop = '__i' + op[2:]
+            ops += [rop, iop]
+        for s in [1, 1., False]:
+            for _op in ops:
+                for a in _array_vals():
+                    # Test array op scalar. From the spec, the following combinations
+                    # are supported:
+
+                    # - Python bool for a bool array dtype,
+                    # - a Python int within the bounds of the given dtype for integer array dtypes,
+                    # - a Python int or float for floating-point array dtypes
+
+                    # We do not do bounds checking for int scalars, but rather use the default
+                    # NumPy behavior for casting in that case.
+
+                    if ((dtypes == "all"
+                         or dtypes == "numeric" and a.dtype in _numeric_dtypes
+                         or dtypes == "integer" and a.dtype in _integer_dtypes
+                         or dtypes == "integer_or_boolean" and a.dtype in _integer_or_boolean_dtypes
+                         or dtypes == "boolean" and a.dtype in _boolean_dtypes
+                         or dtypes == "floating" and a.dtype in _floating_dtypes
+                        )
+                        # bool is a subtype of int, which is why we avoid
+                        # isinstance here.
+                        and (a.dtype in _boolean_dtypes and type(s) == bool
+                             or a.dtype in _integer_dtypes and type(s) == int
+                             or a.dtype in _floating_dtypes and type(s) in [float, int]
+                        )):
+                        # Only test for no error
+                        getattr(a, _op)(s)
+                    else:
+                        assert_raises(TypeError, lambda: getattr(a, _op)(s))
+
+                # Test array op array.
+                for _op in ops:
+                    for x in _array_vals():
+                        for y in _array_vals():
+                            # See the promotion table in NEP 47 or the array
+                            # API spec page on type promotion. Mixed kind
+                            # promotion is not defined.
+                            if (x.dtype == uint64 and y.dtype in [int8, int16, int32, int64]
+                                or y.dtype == uint64 and x.dtype in [int8, int16, int32, int64]
+                                or x.dtype in _integer_dtypes and y.dtype not in _integer_dtypes
+                                or y.dtype in _integer_dtypes and x.dtype not in _integer_dtypes
+                                or x.dtype in _boolean_dtypes and y.dtype not in _boolean_dtypes
+                                or y.dtype in _boolean_dtypes and x.dtype not in _boolean_dtypes
+                                or x.dtype in _floating_dtypes and y.dtype not in _floating_dtypes
+                                or y.dtype in _floating_dtypes and x.dtype not in _floating_dtypes
+                                ):
+                                assert_raises(TypeError, lambda: getattr(x, _op)(y))
+                            # Ensure in-place operators only promote to the same dtype as the left operand.
+                            elif _op.startswith('__i') and result_type(x.dtype, y.dtype) != x.dtype:
+                                assert_raises(TypeError, lambda: getattr(x, _op)(y))
+                            # Ensure only those dtypes that are required for every operator are allowed.
+                            elif (dtypes == "all" and (x.dtype in _boolean_dtypes and y.dtype in _boolean_dtypes
+                                                      or x.dtype in _numeric_dtypes and y.dtype in _numeric_dtypes)
+                                or (dtypes == "numeric" and x.dtype in _numeric_dtypes and y.dtype in _numeric_dtypes)
+                                or dtypes == "integer" and x.dtype in _integer_dtypes and y.dtype in _numeric_dtypes
+                                or dtypes == "integer_or_boolean" and (x.dtype in _integer_dtypes and y.dtype in _integer_dtypes
+                                                                       or x.dtype in _boolean_dtypes and y.dtype in _boolean_dtypes)
+                                or dtypes == "boolean" and x.dtype in _boolean_dtypes and y.dtype in _boolean_dtypes
+                                or dtypes == "floating" and x.dtype in _floating_dtypes and y.dtype in _floating_dtypes
+                            ):
+                                getattr(x, _op)(y)
+                            else:
+                                assert_raises(TypeError, lambda: getattr(x, _op)(y))
+
+    unary_op_dtypes ={
+        '__abs__': 'numeric',
+        '__invert__': 'integer_or_boolean',
+        '__neg__': 'numeric',
+        '__pos__': 'numeric',
+    }
+    for op, dtypes in unary_op_dtypes.items():
+        for a in _array_vals():
+            if (dtypes == "numeric" and a.dtype in _numeric_dtypes
+                or dtypes == "integer_or_boolean" and a.dtype in _integer_or_boolean_dtypes
+                ):
+                # Only test for no error
+                getattr(a, op)()
+            else:
+                assert_raises(TypeError, lambda: getattr(a, op)())
+
+    # Finally, matmul() must be tested separately, because it works a bit
+    # different from the other operations.
+    def _matmul_array_vals():
+        for a in _array_vals():
+            yield a
+        for d in _all_dtypes:
+            yield ones((3, 4), dtype=d)
+            yield ones((4, 2), dtype=d)
+            yield ones((4, 4), dtype=d)
+
+    # Scalars always error
+    for _op in ['__matmul__', '__rmatmul__', '__imatmul__']:
+        for s in [1, 1., False]:
+            for a in _matmul_array_vals():
+                if (type(s) in [float, int] and a.dtype in _floating_dtypes
+                    or type(s) == int and a.dtype in _integer_dtypes):
+                    # Type promotion is valid, but @ is not allowed on 0-D
+                    # inputs, so the error is a ValueError
+                    assert_raises(ValueError, lambda: getattr(a, _op)(s))
+                else:
+                    assert_raises(TypeError, lambda: getattr(a, _op)(s))
+
+    for x in _matmul_array_vals():
+        for y in _matmul_array_vals():
+            if (x.dtype == uint64 and y.dtype in [int8, int16, int32, int64]
+                or y.dtype == uint64 and x.dtype in [int8, int16, int32, int64]
+                or x.dtype in _integer_dtypes and y.dtype not in _integer_dtypes
+                or y.dtype in _integer_dtypes and x.dtype not in _integer_dtypes
+                or x.dtype in _floating_dtypes and y.dtype not in _floating_dtypes
+                or y.dtype in _floating_dtypes and x.dtype not in _floating_dtypes
+                or x.dtype in _boolean_dtypes
+                or y.dtype in _boolean_dtypes
+                ):
+                assert_raises(TypeError, lambda: x.__matmul__(y))
+                assert_raises(TypeError, lambda: y.__rmatmul__(x))
+                assert_raises(TypeError, lambda: x.__imatmul__(y))
+            elif x.shape == () or y.shape == () or x.shape[1] != y.shape[0]:
+                assert_raises(ValueError, lambda: x.__matmul__(y))
+                assert_raises(ValueError, lambda: y.__rmatmul__(x))
+                if result_type(x.dtype, y.dtype) != x.dtype:
+                    assert_raises(TypeError, lambda: x.__imatmul__(y))
+                else:
+                    assert_raises(ValueError, lambda: x.__imatmul__(y))
+            else:
+                x.__matmul__(y)
+                y.__rmatmul__(x)
+                if result_type(x.dtype, y.dtype) != x.dtype:
+                    assert_raises(TypeError, lambda: x.__imatmul__(y))
+                elif y.shape[0] != y.shape[1]:
+                    # This one fails because x @ y has a different shape from x
+                    assert_raises(ValueError, lambda: x.__imatmul__(y))
+                else:
+                    x.__imatmul__(y)
+
+def test_python_scalar_construtors():
+    a = asarray(False)
+    b = asarray(0)
+    c = asarray(0.)
+
+    assert bool(a) == bool(b) == bool(c) == False
+    assert int(a) == int(b) == int(c) == 0
+    assert float(a) == float(b) == float(c) == 0.
+
+    # bool/int/float should only be allowed on 0-D arrays.
+    assert_raises(TypeError, lambda: bool(asarray([False])))
+    assert_raises(TypeError, lambda: int(asarray([0])))
+    assert_raises(TypeError, lambda: float(asarray([0.])))