Add tests for linspace()

Author: asmeurer

array_api_tests/test_creation_functions.py

@@ -1,12 +1,13 @@
 from ._array_module import (arange, ceil, empty, _floating_dtypes, eye, full,
-equal, all)
-from .array_helpers import is_integer_dtype, dtype_ranges
+equal, all, linspace)
+from .array_helpers import (is_integer_dtype, dtype_ranges,
+                            assert_exactly_equal, isintegral)
 from .hypothesis_helpers import (numeric_dtypes, dtypes, MAX_ARRAY_SIZE,
                                  shapes, sizes, sqrt_sizes, shared_dtypes,
                                  shared_scalars)
 
 from hypothesis import assume, given
-from hypothesis.strategies import integers, floats, one_of, none
+from hypothesis.strategies import integers, floats, one_of, none, booleans
 
 int_range = integers(-MAX_ARRAY_SIZE, MAX_ARRAY_SIZE)
 float_range = floats(-MAX_ARRAY_SIZE, MAX_ARRAY_SIZE, allow_nan=False)
@@ -86,7 +87,7 @@ def test_eye(N, M, k, dtype):
     if dtype is None:
         assert a.dtype in _floating_dtypes, "eye() should produce an array with the default floating point dtype"
     else:
-        assert a.dtype == dtype
+        assert a.dtype == dtype, "eye() did not produce the correct dtype"
 
     if M is None:
         M = N
@@ -116,3 +117,55 @@ def test_full(shape, fill_value, dtype):
 
     assert a.shape == shape, "full() produced an array with incorrect shape"
     assert all(equal(a, fill_value)), "full() array did not equal the fill value"
+
+# TODO: implement full_like (requires hypothesis arrays support)
+def test_full_like():
+    pass
+
+@given(one_of(integers(), floats(allow_nan=False, allow_infinity=False)),
+       one_of(integers(), floats(allow_nan=False, allow_infinity=False)),
+       sizes,
+       one_of(none(), dtypes),
+       one_of(none(), booleans()),)
+def test_linspace(start, stop, num, dtype, endpoint):
+    if dtype in dtype_ranges:
+        m, M = dtype_ranges[dtype]
+        if (isinstance(start, int) and not (m <= start <= M)
+            or isinstance(stop, int) and not (m <= stop <= M)):
+            assume(False)
+    # Skip on int start or stop that cannot be exactly represented as a float,
+    # since we do not have good approx_equal helpers yet.
+    if (dtype is None or dtype in _floating_dtypes
+        and ((isinstance(start, int) and not isintegral(start))
+             or (isinstance(stop, int) and not isintegral(stop)))):
+        assume(False)
+
+    kwargs = {k: v for k, v in {'dtype': dtype, 'endpoint': endpoint}.items()
+              if v is not None}
+    a = linspace(start, stop, num, **kwargs)
+
+    if dtype is None:
+        assert a.dtype in _floating_dtypes, "linspace() should produce an array with the default floating point dtype"
+    else:
+        assert a.dtype == dtype, "linspace() did not produce the correct dtype"
+
+    assert a.shape == (num,), "linspace() did not produce an array with the correct shape"
+
+    if endpoint in [None, True]:
+        if num > 1:
+            assert all(equal(a[-1], full((), stop, dtype=dtype))), "linspace() produced an array that does not the endpoint"
+    else:
+        # linspace(..., num, endpoint=False) is the same as the first num
+        # elements of linspace(..., num+1, endpoint=True)
+        b = linspace(start, stop, num + 1, **{**kwargs, 'endpoint': True})
+        assert_exactly_equal(b[:-1], a)
+
+    if num > 0:
+        # We need to cast start to dtype
+        assert all(equal(a[0], full((), start, dtype=dtype))), "linspace() produced an array that does not start with the start"
+
+        # TODO: This requires an assert_approx_equal function
+
+        # n = num - 1 if endpoint in [None, True] else num
+        # for i in range(1, num):
+        #     assert all(equal(a[i], full((), i*(stop - start)/n + start, dtype=dtype))), f"linspace() produced an array with an incorrect value at index {i}"