Remove array helper imports which are covered by xp

Author: honno

array_api_tests/test_elementwise_functions.py

@@ -14,36 +14,31 @@
 
 """
 
-from hypothesis import given, assume
-from hypothesis.strategies import composite, just
-
 import math
 
-from .hypothesis_helpers import (integer_dtype_objects,
-                                 floating_dtype_objects,
-                                 numeric_dtype_objects,
+from hypothesis import assume, given
+from hypothesis import strategies as st
+
+from . import _array_module as xp
+from .array_helpers import (assert_exactly_equal, assert_integral,
+                            assert_same_sign, dtype_ranges, false, infinity,
+                            inrange, int_to_dtype, is_float_dtype,
+                            is_integer_dtype, isintegral, isnegative, ndindex,
+                            negative_mathematical_sign, one,
+                            positive_mathematical_sign, promote_dtypes, true,
+                            zero, π)
+from .hypothesis_helpers import (array_scalars, boolean_dtype_objects,
+                                 boolean_dtypes, floating_dtype_objects,
+                                 floating_dtypes, integer_dtype_objects,
                                  integer_or_boolean_dtype_objects,
-                                 boolean_dtype_objects, floating_dtypes,
-                                 numeric_dtypes, integer_or_boolean_dtypes,
-                                 boolean_dtypes, mutually_promotable_dtypes,
-                                 array_scalars, two_mutual_arrays, xps, shapes)
-from .array_helpers import (assert_exactly_equal, negative,
-                            positive_mathematical_sign,
-                            negative_mathematical_sign, logical_not,
-                            logical_or, logical_and, inrange, π, one, zero,
-                            infinity, isnegative, all as array_all, any as
-                            array_any, int_to_dtype,
-                            assert_integral, less_equal, isintegral, isfinite,
-                            ndindex, promote_dtypes, is_integer_dtype,
-                            is_float_dtype, not_equal, asarray,
-                            dtype_ranges, full, true, false, assert_same_sign,
-                            isnan, equal, less)
+                                 integer_or_boolean_dtypes,
+                                 mutually_promotable_dtypes,
+                                 numeric_dtype_objects, numeric_dtypes, shapes,
+                                 two_mutual_arrays, xps)
 # We might as well use this implementation rather than requiring
 # mod.broadcast_shapes(). See test_equal() and others.
 from .test_broadcasting import broadcast_shapes
 
-from . import _array_module as xp
-
 # integer_scalars = array_scalars(integer_dtypes)
 floating_scalars = array_scalars(floating_dtypes)
 numeric_scalars = array_scalars(numeric_dtypes)
@@ -57,11 +52,11 @@
 two_boolean_dtypes = mutually_promotable_dtypes(boolean_dtype_objects)
 two_any_dtypes = mutually_promotable_dtypes()
 
-@composite
+@st.composite
 def two_array_scalars(draw, dtype1, dtype2):
     # two_dtypes should be a strategy that returns two dtypes (like
     # mutually_promotable_dtypes())
-    return draw(array_scalars(just(dtype1))), draw(array_scalars(just(dtype2)))
+    return draw(array_scalars(st.just(dtype1))), draw(array_scalars(st.just(dtype2)))
 
 def sanity_check(x1, x2):
     try:
@@ -74,17 +69,17 @@ def test_abs(x):
     if is_integer_dtype(x.dtype):
         minval = dtype_ranges[x.dtype][0]
         if minval < 0:
-            # abs of the smallest representable negative integer is not defined
-            mask = not_equal(x, full(x.shape, minval, dtype=x.dtype))
+            # abs of the smallest representable xp.negative integer is not defined
+            mask = xp.not_equal(x, xp.full(x.shape, minval, dtype=x.dtype))
             x = x[mask]
     a = xp.abs(x)
-    assert array_all(logical_not(negative_mathematical_sign(a))), "abs(x) did not have positive sign"
+    assert xp.all(xp.logical_not(negative_mathematical_sign(a))), "abs(x) did not have positive sign"
     less_zero = negative_mathematical_sign(x)
-    negx = negative(x)
+    negx = xp.negative(x)
     # abs(x) = -x for x < 0
     assert_exactly_equal(a[less_zero], negx[less_zero])
     # abs(x) = x for x >= 0
-    assert_exactly_equal(a[logical_not(less_zero)], x[logical_not(less_zero)])
+    assert_exactly_equal(a[xp.logical_not(less_zero)], x[xp.logical_not(less_zero)])
 
 @given(xps.arrays(dtype=xps.floating_dtypes(), shape=shapes))
 def test_acos(x):
@@ -170,7 +165,7 @@ def test_atan2(x1_and_x2):
     # atan2 maps [-inf, inf] x [-inf, inf] to [-pi, pi]. Values outside
     # this domain are mapped to nan, which is already tested in the special
     # cases.
-    assert_exactly_equal(logical_and(domainx1, domainx2), codomain)
+    assert_exactly_equal(xp.logical_and(domainx1, domainx2), codomain)
     # From the spec:
     #
     # The mathematical signs of `x1_i` and `x2_i` determine the quadrant of
@@ -187,10 +182,10 @@ def test_atan2(x1_and_x2):
     negx2 = negative_mathematical_sign(x2)
     posa = positive_mathematical_sign(a)
     nega = negative_mathematical_sign(a)
-    assert_exactly_equal(logical_or(logical_and(posx1, posx2),
-                                    logical_and(posx1, negx2)), posa)
-    assert_exactly_equal(logical_or(logical_and(negx1, posx2),
-                                    logical_and(negx1, negx2)), nega)
+    assert_exactly_equal(xp.logical_or(xp.logical_and(posx1, posx2),
+                                    xp.logical_and(posx1, negx2)), posa)
+    assert_exactly_equal(xp.logical_or(xp.logical_and(negx1, posx2),
+                                    xp.logical_and(negx1, negx2)), nega)
 
 @given(xps.arrays(dtype=xps.floating_dtypes(), shape=shapes))
 def test_atanh(x):
@@ -232,8 +227,7 @@ def test_bitwise_left_shift(args):
     x1, x2 = args
     sanity_check(x1, x2)
     negative_x2 = isnegative(x2)
-    if array_any(negative_x2):
-        assume(False)
+    assume(not xp.any(negative_x2))
     a = xp.bitwise_left_shift(x1, x2)
     # Compare against the Python << operator.
     # TODO: Generalize this properly for inputs that are arrays.
@@ -296,8 +290,7 @@ def test_bitwise_right_shift(args):
     x1, x2 = args
     sanity_check(x1, x2)
     negative_x2 = isnegative(x2)
-    if array_any(negative_x2):
-        assume(False)
+    assume(not xp.any(negative_x2))
     a = xp.bitwise_right_shift(x1, x2)
     # Compare against the Python >> operator.
     # TODO: Generalize this properly for inputs that are arrays.
@@ -335,10 +328,10 @@ def test_bitwise_xor(args):
 def test_ceil(x):
     # This test is almost identical to test_floor()
     a = xp.ceil(x)
-    finite = isfinite(x)
+    finite = xp.isfinite(x)
     assert_integral(a[finite])
-    assert array_all(less_equal(x[finite], a[finite]))
-    assert array_all(less_equal(a[finite] - x[finite], one(x[finite].shape, x.dtype)))
+    assert xp.all(xp.less_equal(x[finite], a[finite]))
+    assert xp.all(xp.less_equal(a[finite] - x[finite], one(x[finite].shape, x.dtype)))
     integers = isintegral(x)
     assert_exactly_equal(a[integers], x[integers])
 
@@ -380,7 +373,7 @@ def test_equal(x1_and_x2):
     x1, x2 = x1_and_x2
     sanity_check(x1, x2)
     a = xp.equal(x1, x2)
-    # NOTE: assert_exactly_equal() itself uses equal(), so we must be careful
+    # NOTE: assert_exactly_equal() itself uses xp.equal(), so we must be careful
     # not to use it here. Otherwise, the test would be circular and
     # meaningless. Instead, we implement this by iterating every element of
     # the arrays and comparing them. The logic here is also used for the tests
@@ -397,15 +390,15 @@ def test_equal(x1_and_x2):
     _x2 = xp.broadcast_to(x2, shape)
 
     # Second, manually promote the dtypes. This is important. If the internal
-    # type promotion in equal() is wrong, it will not be directly visible in
+    # type promotion in xp.equal() is wrong, it will not be directly visible in
     # the output type, but it can lead to wrong answers. For example,
-    # equal(array(1.0, dtype=float32), array(1.00000001, dtype=xp.float64)) will
-    # be wrong if the xp.float64 is downcast to float32. # be wrong if the
+    # xp.equal(array(1.0, dtype=xp.float32), array(1.00000001, dtype=xp.float64)) will
+    # be wrong if the float64 is downcast to float32. # be wrong if the
     # xp.float64 is downcast to float32. See the comment on
     # test_elementwise_function_two_arg_bool_type_promotion() in
-    # test_type_promotion.py. The type promotion for equal() is not *really*
+    # test_type_promotion.py. The type promotion for xp.equal() is not *really*
     # tested in that file, because doing so requires doing the consistency
-    # check we do here rather than just checking the result dtype.
+    # check we do here rather than st.just checking the result dtype.
     promoted_dtype = promote_dtypes(x1.dtype, x2.dtype)
     _x1 = xp.asarray(_x1, dtype=promoted_dtype)
     _x2 = xp.asarray(_x2, dtype=promoted_dtype)
@@ -450,10 +443,10 @@ def test_expm1(x):
 def test_floor(x):
     # This test is almost identical to test_ceil
     a = xp.floor(x)
-    finite = isfinite(x)
+    finite = xp.isfinite(x)
     assert_integral(a[finite])
-    assert array_all(less_equal(a[finite], x[finite]))
-    assert array_all(less_equal(x[finite] - a[finite], one(x[finite].shape, x.dtype)))
+    assert xp.all(xp.less_equal(a[finite], x[finite]))
+    assert xp.all(xp.less_equal(x[finite] - a[finite], one(x[finite].shape, x.dtype)))
     integers = isintegral(x)
     assert_exactly_equal(a[integers], x[integers])
 
@@ -467,8 +460,8 @@ def test_floor_divide(x1_and_x2):
         # we avoid passing it in entirely.
         assume(not xp.any(x1 == 0) and not xp.any(x2 == 0))
         div = xp.divide(
-            asarray(x1, dtype=xp.float64),
-            asarray(x2, dtype=xp.float64),
+            xp.asarray(x1, dtype=xp.float64),
+            xp.asarray(x2, dtype=xp.float64),
         )
     else:
         div = xp.divide(x1, x2)
@@ -477,7 +470,7 @@ def test_floor_divide(x1_and_x2):
 
     # TODO: The spec doesn't clearly specify the behavior of floor_divide on
     # infinities. See https://github.com/data-apis/array-api/issues/199.
-    finite = isfinite(div)
+    finite = xp.isfinite(div)
     assert_integral(out[finite])
 
     # TODO: Test the exact output for floor_divide.
@@ -548,9 +541,9 @@ def test_isfinite(x):
     TRUE = true(x.shape)
     if is_integer_dtype(x.dtype):
         assert_exactly_equal(a, TRUE)
-    # Test that isfinite, isinf, and isnan are self-consistent.
-    inf = logical_or(xp.isinf(x), xp.isnan(x))
-    assert_exactly_equal(a, logical_not(inf))
+    # Test that xp.isfinite, isinf, and xp.isnan are self-consistent.
+    inf = xp.logical_or(xp.isinf(x), xp.isnan(x))
+    assert_exactly_equal(a, xp.logical_not(inf))
 
     # Test the exact value by comparing to the math version
     if is_float_dtype(x.dtype):
@@ -564,8 +557,8 @@ def test_isinf(x):
     FALSE = false(x.shape)
     if is_integer_dtype(x.dtype):
         assert_exactly_equal(a, FALSE)
-    finite_or_nan = logical_or(xp.isfinite(x), xp.isnan(x))
-    assert_exactly_equal(a, logical_not(finite_or_nan))
+    finite_or_nan = xp.logical_or(xp.isfinite(x), xp.isnan(x))
+    assert_exactly_equal(a, xp.logical_not(finite_or_nan))
 
     # Test the exact value by comparing to the math version
     if is_float_dtype(x.dtype):
@@ -579,8 +572,8 @@ def test_isnan(x):
     FALSE = false(x.shape)
     if is_integer_dtype(x.dtype):
         assert_exactly_equal(a, FALSE)
-    finite_or_inf = logical_or(xp.isfinite(x), xp.isinf(x))
-    assert_exactly_equal(a, logical_not(finite_or_inf))
+    finite_or_inf = xp.logical_or(xp.isfinite(x), xp.isinf(x))
+    assert_exactly_equal(a, xp.logical_not(finite_or_inf))
 
     # Test the exact value by comparing to the math version
     if is_float_dtype(x.dtype):
@@ -767,12 +760,12 @@ def test_negative(x):
     # Negation is an involution
     assert_exactly_equal(x, xp.negative(out))
 
-    mask = isfinite(x)
+    mask = xp.isfinite(x)
     if is_integer_dtype(x.dtype):
         minval = dtype_ranges[x.dtype][0]
         if minval < 0:
             # negative of the smallest representable negative integer is not defined
-            mask = not_equal(x, full(x.shape, minval, dtype=x.dtype))
+            mask = xp.not_equal(x, xp.full(x.shape, minval, dtype=x.dtype))
 
     # Additive inverse
     y = xp.add(x[mask], out[mask])
@@ -837,15 +830,15 @@ def test_remainder(x1_and_x2):
 
     # out and x2 should have the same sign.
     # assert_same_sign returns False for nans
-    not_nan = logical_not(logical_or(isnan(out), isnan(x2)))
+    not_nan = xp.logical_not(xp.logical_or(xp.isnan(out), xp.isnan(x2)))
     assert_same_sign(out[not_nan], x2[not_nan])
 
 @given(xps.arrays(dtype=xps.numeric_dtypes(), shape=shapes))
 def test_round(x):
     a = xp.round(x)
 
     # Test that the result is integral
-    finite = isfinite(x)
+    finite = xp.isfinite(x)
     assert_integral(a[finite])
 
     # round(x) should be the nearest integer to x. The case where there is a
@@ -858,8 +851,8 @@ def test_round(x):
     ceil = xp.ceil(x)
     over = xp.subtract(x, floor)
     under = xp.subtract(ceil, x)
-    round_down = less(over, under)
-    round_up = less(under, over)
+    round_down = xp.less(over, under)
+    round_up = xp.less(under, over)
     assert_exactly_equal(a[round_down], floor[round_down])
     assert_exactly_equal(a[round_up], ceil[round_up])
 
@@ -910,7 +903,7 @@ def test_trunc(x):
     assert out.dtype == x.dtype, f"{x.dtype=!s} but {out.dtype=!s}"
     assert out.shape == x.shape, f"{x.shape} but {out.shape}"
     if x.dtype in integer_dtype_objects:
-        assert array_all(equal(x, out)), f"{x=!s} but {out=!s}"
+        assert xp.all(xp.equal(x, out)), f"{x=!s} but {out=!s}"
     else:
         finite_mask = xp.isfinite(out)
         for idx in ndindex(out.shape):