Merge branch 'master' into type-promotion-refactor

Author: honno

array_api_tests/__init__.py

@@ -0,0 +1,6 @@
+from hypothesis.extra.array_api import make_strategies_namespace
+
+from . import _array_module as xp
+
+
+xps = make_strategies_namespace(xp)

array_api_tests/function_stubs/creation_functions.py

@@ -12,7 +12,6 @@
 
 from ._types import (List, NestedSequence, Optional, SupportsBufferProtocol, SupportsDLPack, Tuple,
                      Union, array, device, dtype)
-from collections.abc import Sequence
 
 def arange(start: Union[int, float], /, stop: Optional[Union[int, float]] = None, step: Union[int, float] = 1, *, dtype: Optional[dtype] = None, device: Optional[device] = None) -> array:
     pass
@@ -26,7 +25,7 @@ def empty(shape: Union[int, Tuple[int, ...]], *, dtype: Optional[dtype] = None,
 def empty_like(x: array, /, *, dtype: Optional[dtype] = None, device: Optional[device] = None) -> array:
     pass
 
-def eye(n_rows: int, n_cols: Optional[int] = None, /, *, k: Optional[int] = 0, dtype: Optional[dtype] = None, device: Optional[device] = None) -> array:
+def eye(n_rows: int, n_cols: Optional[int] = None, /, *, k: int = 0, dtype: Optional[dtype] = None, device: Optional[device] = None) -> array:
     pass
 
 def from_dlpack(x: object, /) -> array:
@@ -41,7 +40,7 @@ def full_like(x: array, /, fill_value: Union[int, float], *, dtype: Optional[dty
 def linspace(start: Union[int, float], stop: Union[int, float], /, num: int, *, dtype: Optional[dtype] = None, device: Optional[device] = None, endpoint: bool = True) -> array:
     pass
 
-def meshgrid(*arrays: Sequence[array], indexing: str = 'xy') -> List[array, ...]:
+def meshgrid(*arrays: array, indexing: str = 'xy') -> List[array, ...]:
     pass
 
 def ones(shape: Union[int, Tuple[int, ...]], *, dtype: Optional[dtype] = None, device: Optional[device] = None) -> array:

array_api_tests/function_stubs/data_type_functions.py

@@ -11,9 +11,8 @@
 from __future__ import annotations
 
 from ._types import List, Tuple, Union, array, dtype, finfo_object, iinfo_object
-from collections.abc import Sequence
 
-def broadcast_arrays(*arrays: Sequence[array]) -> List[array]:
+def broadcast_arrays(*arrays: array) -> List[array]:
     pass
 
 def broadcast_to(x: array, /, shape: Tuple[int, ...]) -> array:
@@ -28,7 +27,7 @@ def finfo(type: Union[dtype, array], /) -> finfo_object:
 def iinfo(type: Union[dtype, array], /) -> iinfo_object:
     pass
 
-def result_type(*arrays_and_dtypes: Sequence[Union[array, dtype]]) -> dtype:
+def result_type(*arrays_and_dtypes: Union[array, dtype]) -> dtype:
     pass
 
 __all__ = ['broadcast_arrays', 'broadcast_to', 'can_cast', 'finfo', 'iinfo', 'result_type']

array_api_tests/function_stubs/linalg.py

@@ -68,7 +68,7 @@ def solve(x1: array, x2: array, /) -> array:
 def svd(x: array, /, *, full_matrices: bool = True) -> Tuple[array, array, array]:
     pass
 
-def svdvals(x: array, /) -> Union[array, Tuple[array, ...]]:
+def svdvals(x: array, /) -> array:
     pass
 
 def tensordot(x1: array, x2: array, /, *, axes: Union[int, Tuple[Sequence[int], Sequence[int]]] = 2) -> array:
@@ -77,10 +77,10 @@ def tensordot(x1: array, x2: array, /, *, axes: Union[int, Tuple[Sequence[int],
 def trace(x: array, /, *, offset: int = 0) -> array:
     pass
 
-def vecdot(x1: array, x2: array, /, *, axis: Optional[int] = None) -> array:
+def vecdot(x1: array, x2: array, /, *, axis: int = None) -> array:
     pass
 
-def vector_norm(x: array, /, *, axis: Optional[Union[int, Tuple[int, int]]] = None, keepdims: bool = False, ord: Optional[Union[int, float, Literal[inf, -inf]]] = 2) -> array:
+def vector_norm(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, keepdims: bool = False, ord: Union[int, float, Literal[inf, -inf]] = 2) -> array:
     pass
 
 __all__ = ['cholesky', 'cross', 'det', 'diagonal', 'eigh', 'eigvalsh', 'inv', 'matmul', 'matrix_norm', 'matrix_power', 'matrix_rank', 'matrix_transpose', 'outer', 'pinv', 'qr', 'slogdet', 'solve', 'svd', 'svdvals', 'tensordot', 'trace', 'vecdot', 'vector_norm']

array_api_tests/function_stubs/linear_algebra_functions.py

@@ -10,7 +10,7 @@
 
 from __future__ import annotations
 
-from ._types import Optional, Tuple, Union, array
+from ._types import Tuple, Union, array
 from collections.abc import Sequence
 
 def matmul(x1: array, x2: array, /) -> array:
@@ -22,7 +22,7 @@ def matrix_transpose(x: array, /) -> array:
 def tensordot(x1: array, x2: array, /, *, axes: Union[int, Tuple[Sequence[int], Sequence[int]]] = 2) -> array:
     pass
 
-def vecdot(x1: array, x2: array, /, *, axis: Optional[int] = None) -> array:
+def vecdot(x1: array, x2: array, /, *, axis: int = -1) -> array:
     pass
 
 __all__ = ['matmul', 'matrix_transpose', 'tensordot', 'vecdot']

array_api_tests/hypothesis_helpers.py

@@ -1,29 +1,28 @@
 from functools import reduce
 from operator import mul
 from math import sqrt
+import itertools
 
 from hypothesis import assume
 from hypothesis.strategies import (lists, integers, sampled_from,
                                    shared, floats, just, composite, one_of,
                                    none, booleans)
-from hypothesis.extra.array_api import make_strategies_namespace
 
 from .pytest_helpers import nargs
 from .array_helpers import (dtype_ranges, integer_dtype_objects,
                             floating_dtype_objects, numeric_dtype_objects,
                             boolean_dtype_objects,
-                            integer_or_boolean_dtype_objects, dtype_objects)
-from ._array_module import full, float32, float64, bool as bool_dtype, _UndefinedStub
+                            integer_or_boolean_dtype_objects, dtype_objects,
+                            ndindex)
 from .dtype_helpers import promotion_table
-from . import _array_module
+from ._array_module import (full, float32, float64, bool as bool_dtype,
+                            _UndefinedStub, eye, broadcast_to)
 from . import _array_module as xp
+from . import xps
 
 from .function_stubs import elementwise_functions
 
 
-xps = make_strategies_namespace(xp)
-
-
 # Set this to True to not fail tests just because a dtype isn't implemented.
 # If no compatible dtype is implemented for a given test, the test will fail
 # with a hypothesis health check error. Note that this functionality will not
@@ -48,6 +47,7 @@
     dtypes = dtypes.filter(lambda x: not isinstance(x, _UndefinedStub))
 
 shared_dtypes = shared(dtypes, key="dtype")
+shared_floating_dtypes = shared(floating_dtypes, key="dtype")
 
 
 sorted_table = sorted(promotion_table)
@@ -72,10 +72,6 @@ def mutually_promotable_dtypes(dtype_objects=dtype_objects):
         [(i, j) for i, j in sorted_table if i in dtype_objects and j in dtype_objects]
     )
 
-shared_mutually_promotable_dtype_pairs = shared(
-    mutually_promotable_dtypes(), key="mutually_promotable_dtype_pair"
-)
-
 # shared() allows us to draw either the function or the function name and they
 # will both correspond to the same function.
 
@@ -86,10 +82,10 @@ def mutually_promotable_dtypes(dtype_objects=dtype_objects):
     lambda func_name: nargs(func_name) > 1)
 
 elementwise_function_objects = elementwise_functions_names.map(
-    lambda i: getattr(_array_module, i))
+    lambda i: getattr(xp, i))
 array_functions = elementwise_function_objects
 multiarg_array_functions = multiarg_array_functions_names.map(
-    lambda i: getattr(_array_module, i))
+    lambda i: getattr(xp, i))
 
 # Limit the total size of an array shape
 MAX_ARRAY_SIZE = 10000
@@ -111,25 +107,72 @@ def tuples(elements, *, min_size=0, max_size=None, unique_by=None, unique=False)
     lambda shape: prod(i for i in shape if i) < MAX_ARRAY_SIZE
 )
 
+# Matrix shapes assume stacks of matrices
+matrix_shapes = xps.array_shapes(min_dims=2, min_side=1).filter(
+    lambda shape: prod(i for i in shape if i) < MAX_ARRAY_SIZE
+)
+
+square_matrix_shapes = matrix_shapes.filter(lambda shape: shape[-1] == shape[-2])
+
 two_mutually_broadcastable_shapes = xps.mutually_broadcastable_shapes(num_shapes=2)\
     .map(lambda S: S.input_shapes)\
     .filter(lambda S: all(prod(i for i in shape if i) < MAX_ARRAY_SIZE for shape in S))
 
+# Note: This should become hermitian_matrices when complex dtypes are added
+@composite
+def symmetric_matrices(draw, dtypes=xps.floating_dtypes(), finite=True):
+    shape = draw(square_matrix_shapes)
+    dtype = draw(dtypes)
+    elements = {'allow_nan': False, 'allow_infinity': False} if finite else None
+    a = draw(xps.arrays(dtype=dtype, shape=shape, elements=elements))
+    upper = xp.triu(a)
+    lower = xp.triu(a, k=1).mT
+    return upper + lower
+
+@composite
+def positive_definite_matrices(draw, dtypes=xps.floating_dtypes()):
+    # For now just generate stacks of identity matrices
+    # TODO: Generate arbitrary positive definite matrices, for instance, by
+    # using something like
+    # https://github.com/scikit-learn/scikit-learn/blob/844b4be24/sklearn/datasets/_samples_generator.py#L1351.
+    n = draw(integers(0))
+    shape = draw(shapes) + (n, n)
+    assume(prod(i for i in shape if i) < MAX_ARRAY_SIZE)
+    dtype = draw(dtypes)
+    return broadcast_to(eye(n, dtype=dtype), shape)
+
+@composite
+def invertible_matrices(draw, dtypes=xps.floating_dtypes()):
+    # For now, just generate stacks of diagonal matrices.
+    n = draw(integers(0, SQRT_MAX_ARRAY_SIZE),)
+    stack_shape = draw(shapes)
+    shape = stack_shape + (n, n)
+    d = draw(xps.arrays(dtypes, shape=n*prod(stack_shape),
+                        elements=dict(allow_nan=False, allow_infinity=False)))
+    # Functions that require invertible matrices may do anything when it is
+    # singular, including raising an exception, so we make sure the diagonals
+    # are sufficiently nonzero to avoid any numerical issues.
+    assume(xp.all(xp.abs(d) > 0.5))
+
+    a = xp.zeros(shape)
+    for j, (idx, i) in enumerate(itertools.product(ndindex(stack_shape), range(n))):
+        a[idx + (i, i)] = d[j]
+    return a
+
 @composite
 def two_broadcastable_shapes(draw):
     """
     This will produce two shapes (shape1, shape2) such that shape2 can be
     broadcast to shape1.
     """
     from .test_broadcasting import broadcast_shapes
-    shape1, shape2 =  draw(two_mutually_broadcastable_shapes)
+    shape1, shape2 = draw(two_mutually_broadcastable_shapes)
     assume(broadcast_shapes(shape1, shape2) == shape1)
     return (shape1, shape2)
 
 sizes = integers(0, MAX_ARRAY_SIZE)
 sqrt_sizes = integers(0, SQRT_MAX_ARRAY_SIZE)
 
-# TODO: Generate general arrays here, rather than just scalars.
 numeric_arrays = xps.arrays(
     dtype=shared(xps.floating_dtypes(), key='dtypes'),
     shape=shared(xps.array_shapes(), key='shapes'),
@@ -233,25 +276,46 @@ def multiaxis_indices(draw, shapes):
 
     # Avoid using 'in', which might do == on an array.
     res_has_ellipsis = any(i is ... for i in res)
-    if n_entries == len(shape) and not res_has_ellipsis:
-        # note("Adding extra")
-        extra = draw(lists(one_of(integer_indices(sizes), slices(sizes)), min_size=0, max_size=3))
-        res += extra
+    if not res_has_ellipsis:
+        if n_entries < len(shape):
+            # The spec requires either an ellipsis or exactly as many indices
+            # as dimensions.
+            assume(False)
+        elif n_entries == len(shape):
+            # note("Adding extra")
+            extra = draw(lists(one_of(integer_indices(sizes), slices(sizes)), min_size=0, max_size=3))
+            res += extra
     return tuple(res)
 
 
-shared_arrays1 = xps.arrays(
-    dtype=shared_mutually_promotable_dtype_pairs.map(lambda pair: pair[0]),
-    shape=shared(two_mutually_broadcastable_shapes, key="shape_pair").map(lambda pair: pair[0]),
-)
-shared_arrays2 = xps.arrays(
-    dtype=shared_mutually_promotable_dtype_pairs.map(lambda pair: pair[1]),
-    shape=shared(two_mutually_broadcastable_shapes, key="shape_pair").map(lambda pair: pair[1]),
-)
+def two_mutual_arrays(dtype_objects=dtype_objects):
+    mutual_dtypes = shared(mutually_promotable_dtypes(dtype_objects))
+    mutual_shapes = shared(two_mutually_broadcastable_shapes)
+    arrays1 = xps.arrays(
+        dtype=mutual_dtypes.map(lambda pair: pair[0]),
+        shape=mutual_shapes.map(lambda pair: pair[0]),
+    )
+    arrays2 = xps.arrays(
+        dtype=mutual_dtypes.map(lambda pair: pair[1]),
+        shape=mutual_shapes.map(lambda pair: pair[1]),
+    )
+    return arrays1, arrays2
 
 
 @composite
 def kwargs(draw, **kw):
+    """
+    Strategy for keyword arguments
+
+    For a signature like f(x, /, dtype=None, val=1) use
+
+    @given(x=arrays(), kw=kwargs(a=none() | dtypes, val=integers()))
+    def test_f(x, kw):
+        res = f(x, **kw)
+
+    kw may omit the keyword argument, meaning the default for f will be used.
+
+    """
     result = {}
     for k, strat in kw.items():
         if draw(booleans()):

array_api_tests/meta_tests/test_hypothesis_helpers.py

@@ -1,17 +1,18 @@
 from math import prod
 
 import pytest
-from hypothesis import given, strategies as st
+from hypothesis import given, strategies as st, settings
 
 from .. import _array_module as xp
 from .._array_module import _UndefinedStub
 from .. import array_helpers as ah
 from .. import hypothesis_helpers as hh
+from ..test_broadcasting import broadcast_shapes
+from ..test_elementwise_functions import sanity_check
 
 UNDEFINED_DTYPES = any(isinstance(d, _UndefinedStub) for d in ah.dtype_objects)
 pytestmark = [pytest.mark.skipif(UNDEFINED_DTYPES, reason="undefined dtypes")]
 
-
 @given(hh.mutually_promotable_dtypes([xp.float32, xp.float64]))
 def test_mutually_promotable_dtypes(pairs):
     assert pairs in (
@@ -45,20 +46,24 @@ def test_two_mutually_broadcastable_shapes(pair):
 def test_two_broadcastable_shapes(pair):
     for shape in pair:
         assert valid_shape(shape)
+    assert broadcast_shapes(pair[0], pair[1]) == pair[0]
 
-    from ..test_broadcasting import broadcast_shapes
 
-    assert broadcast_shapes(pair[0], pair[1]) == pair[0]
+@given(*hh.two_mutual_arrays())
+def test_two_mutual_arrays(x1, x2):
+    sanity_check(x1, x2)
+    assert broadcast_shapes(x1.shape, x2.shape) in (x1.shape, x2.shape)
 
 
 def test_kwargs():
     results = []
 
     @given(hh.kwargs(n=st.integers(0, 10), c=st.from_regex("[a-f]")))
+    @settings(max_examples=100)
     def run(kw):
         results.append(kw)
-
     run()
+
     assert all(isinstance(kw, dict) for kw in results)
     for size in [0, 1, 2]:
         assert any(len(kw) == size for kw in results)
@@ -70,3 +75,21 @@ def run(kw):
     c_results = [kw for kw in results if "c" in kw]
     assert len(c_results) > 0
     assert all(isinstance(kw["c"], str) for kw in c_results)
+
+@given(m=hh.symmetric_matrices(hh.shared_floating_dtypes,
+                                     finite=st.shared(st.booleans(), key='finite')),
+       dtype=hh.shared_floating_dtypes,
+       finite=st.shared(st.booleans(), key='finite'))
+def test_symmetric_matrices(m, dtype, finite):
+    assert m.dtype == dtype
+    # TODO: This part of this test should be part of the .mT test
+    ah.assert_exactly_equal(m, m.mT)
+
+    if finite:
+        ah.assert_finite(m)
+
+@given(m=hh.positive_definite_matrices(hh.shared_floating_dtypes),
+       dtype=hh.shared_floating_dtypes)
+def test_positive_definite_matrices(m, dtype):
+    assert m.dtype == dtype
+    # TODO: Test that it actually is positive definite

array_api_tests/test_creation_functions.py

@@ -8,7 +8,8 @@
                             assert_exactly_equal, isintegral, is_float_dtype)
 from .hypothesis_helpers import (numeric_dtypes, dtypes, MAX_ARRAY_SIZE,
                                  shapes, sizes, sqrt_sizes, shared_dtypes,
-                                 scalars, xps, kwargs)
+                                 scalars, kwargs)
+from . import xps
 
 from hypothesis import assume, given
 from hypothesis.strategies import integers, floats, one_of, none, booleans, just, shared, composite