Merge pull request data-apis#371 from ev-br/test_pow_with_scalars

test pow() and bitwise shifts with scalars

Author: ev-br

array_api_tests/hypothesis_helpers.py

@@ -456,8 +456,12 @@ def scalars(draw, dtypes, finite=False, **kwds):
     dtypes should be one of the shared_* dtypes strategies.
     """
     dtype = draw(dtypes)
+    mM = kwds.pop('mM', None)
     if dh.is_int_dtype(dtype):
-        m, M = dh.dtype_ranges[dtype]
+        if mM is None:
+            m, M = dh.dtype_ranges[dtype]
+        else:
+            m, M = mM
         return draw(integers(m, M))
     elif dtype == bool_dtype:
         return draw(booleans())
@@ -588,18 +592,20 @@ def two_mutual_arrays(
 
 
 @composite
-def array_and_py_scalar(draw, dtypes):
+def array_and_py_scalar(draw, dtypes, mM=None, positive=False):
     """Draw a pair: (array, scalar) or (scalar, array)."""
     dtype = draw(sampled_from(dtypes))
 
-    scalar_var = draw(scalars(just(dtype), finite=True,
-        **{'min_value': 1/ (2<<5), 'max_value': 2<<5}
-    ))
+    scalar_var = draw(scalars(just(dtype), finite=True, mM=mM))
+    if positive:
+        assume (scalar_var > 0)
 
     elements={}
     if dtype in dh.real_float_dtypes:
         elements = {'allow_nan': False, 'allow_infinity': False,
                     'min_value': 1.0 / (2<<5), 'max_value': 2<<5}
+    if positive:
+        elements = {'min_value': 0}
     array_var = draw(arrays(dtype, shape=shapes(min_dims=1), elements=elements))
 
     if draw(booleans()):

array_api_tests/test_operators_and_elementwise_functions.py

@@ -724,7 +724,9 @@ def _assert_correctness_binary(
     x1a, x2a = in_arrs
     ph.assert_dtype(name, in_dtype=in_dtypes, out_dtype=out.dtype, expected=expected_dtype)
     ph.assert_result_shape(name, in_shapes=in_shapes, out_shape=out.shape)
-    binary_assert_against_refimpl(name, x1a, x2a, out, func, **kwargs)
+    check_values = kwargs.pop('check_values', None)
+    if check_values:
+        binary_assert_against_refimpl(name, x1a, x2a, out, func, **kwargs)
 
 
 @pytest.mark.parametrize("ctx", make_unary_params("abs", dh.numeric_dtypes))
@@ -1845,6 +1847,7 @@ def _filter_zero(x):
         ("less_equal", operator.le, {}, xp.bool),
         ("greater", operator.gt, {}, xp.bool),
         ("greater_equal", operator.ge, {}, xp.bool),
+        ("pow", operator.pow, {'check_values': False}, None)   # value tests are too finicky for pow
     ],
     ids=lambda func_data: func_data[0]  # use names for test IDs
 )
@@ -1902,6 +1905,23 @@ def test_binary_with_scalars_bitwise(func_data, x1x2):
     _check_binary_with_scalars((func_name, refimpl_, kwargs, expected), x1x2)
 
 
+@pytest.mark.min_version("2024.12")
+@pytest.mark.parametrize('func_data',
+    # func_name, refimpl, kwargs, expected_dtype
+    [
+        ("bitwise_left_shift", operator.lshift, {}, None),
+        ("bitwise_right_shift", operator.rshift, {}, None),
+    ],
+    ids=lambda func_data: func_data[0]  # use names for test IDs
+)
+@given(x1x2=hh.array_and_py_scalar([xp.int32], positive=True, mM=(1, 3)))
+def test_binary_with_scalars_bitwise_shifts(func_data, x1x2):
+    func_name, refimpl, kwargs, expected = func_data
+    # repack the refimpl
+    refimpl_ = lambda l, r: mock_int_dtype(refimpl(l, r), xp.int32 )
+    _check_binary_with_scalars((func_name, refimpl_, kwargs, expected), x1x2)
+
+
 @pytest.mark.unvectorized
 @given(
     x1x2=hh.array_and_py_scalar([xp.int32]),
@@ -1931,3 +1951,4 @@ def test_where_with_scalars(x1x2, data):
         else:
             assert out[idx] == x2_arr[idx]
 
+