Merge pull request #82 from honno/dtype-funcs

Test dtype functions

Author: honno

array_api_tests/dtype_helpers.py

@@ -86,8 +86,8 @@ def get_scalar_type(dtype: DataType) -> ScalarType:
 
 
 class MinMax(NamedTuple):
-    min: int
-    max: int
+    min: Union[int, float]
+    max: Union[int, float]
 
 
 dtype_ranges = {
@@ -99,6 +99,8 @@ class MinMax(NamedTuple):
     xp.uint16: MinMax(0, +65_535),
     xp.uint32: MinMax(0, +4_294_967_295),
     xp.uint64: MinMax(0, +18_446_744_073_709_551_615),
+    xp.float32: MinMax(-3.4028234663852886e+38, 3.4028234663852886e+38),
+    xp.float64: MinMax(-1.7976931348623157e+308, 1.7976931348623157e+308),
 }
 
 dtype_nbits = {

array_api_tests/test_data_type_functions.py

@@ -1,12 +1,131 @@
+import struct
+from typing import Union
+
 import pytest
 from hypothesis import given
+from hypothesis import strategies as st
 
 from . import _array_module as xp
 from . import dtype_helpers as dh
 from . import hypothesis_helpers as hh
 from . import pytest_helpers as ph
+from . import xps
+from .algos import broadcast_shapes
 from .typing import DataType
 
+pytestmark = pytest.mark.ci
+
+
+def float32(n: Union[int, float]) -> float:
+    return struct.unpack("!f", struct.pack("!f", float(n)))[0]
+
+
+@given(
+    x_dtype=xps.scalar_dtypes(),
+    dtype=xps.scalar_dtypes(),
+    kw=hh.kwargs(copy=st.booleans()),
+    data=st.data(),
+)
+def test_astype(x_dtype, dtype, kw, data):
+    if xp.bool in (x_dtype, dtype):
+        elements_strat = xps.from_dtype(x_dtype)
+    else:
+        m1, M1 = dh.dtype_ranges[x_dtype]
+        m2, M2 = dh.dtype_ranges[dtype]
+        if dh.is_int_dtype(x_dtype):
+            cast = int
+        elif x_dtype == xp.float32:
+            cast = float32
+        else:
+            cast = float
+        min_value = cast(max(m1, m2))
+        max_value = cast(min(M1, M2))
+        elements_strat = xps.from_dtype(
+            x_dtype,
+            min_value=min_value,
+            max_value=max_value,
+            allow_nan=False,
+            allow_infinity=False,
+        )
+    x = data.draw(
+        xps.arrays(dtype=x_dtype, shape=hh.shapes(), elements=elements_strat), label="x"
+    )
+
+    out = xp.astype(x, dtype, **kw)
+
+    ph.assert_kw_dtype("astype", dtype, out.dtype)
+    ph.assert_shape("astype", out.shape, x.shape)
+    # TODO: test values
+    # TODO: test copy
+
+
+@given(
+    shapes=st.integers(1, 5).flatmap(hh.mutually_broadcastable_shapes), data=st.data()
+)
+def test_broadcast_arrays(shapes, data):
+    arrays = []
+    for c, shape in enumerate(shapes, 1):
+        x = data.draw(xps.arrays(dtype=xps.scalar_dtypes(), shape=shape), label=f"x{c}")
+        arrays.append(x)
+
+    out = xp.broadcast_arrays(*arrays)
+
+    out_shape = broadcast_shapes(*shapes)
+    for i, x in enumerate(arrays):
+        ph.assert_dtype(
+            "broadcast_arrays", x.dtype, out[i].dtype, repr_name=f"out[{i}].dtype"
+        )
+        ph.assert_result_shape(
+            "broadcast_arrays",
+            shapes,
+            out[i].shape,
+            out_shape,
+            repr_name=f"out[{i}].shape",
+        )
+    # TODO: test values
+
+
+@given(x=xps.arrays(dtype=xps.scalar_dtypes(), shape=hh.shapes()), data=st.data())
+def test_broadcast_to(x, data):
+    shape = data.draw(
+        hh.mutually_broadcastable_shapes(1, base_shape=x.shape)
+        .map(lambda S: S[0])
+        .filter(lambda s: broadcast_shapes(x.shape, s) == s),
+        label="shape",
+    )
+
+    out = xp.broadcast_to(x, shape)
+
+    ph.assert_dtype("broadcast_to", x.dtype, out.dtype)
+    ph.assert_shape("broadcast_to", out.shape, shape)
+    # TODO: test values
+
+
+@given(_from=xps.scalar_dtypes(), to=xps.scalar_dtypes(), data=st.data())
+def test_can_cast(_from, to, data):
+    from_ = data.draw(
+        st.just(_from) | xps.arrays(dtype=_from, shape=hh.shapes()), label="from_"
+    )
+
+    out = xp.can_cast(from_, to)
+
+    f_func = f"[can_cast({dh.dtype_to_name[_from]}, {dh.dtype_to_name[to]})]"
+    assert isinstance(out, bool), f"{type(out)=}, but should be bool {f_func}"
+    if _from == xp.bool:
+        expected = to == xp.bool
+    else:
+        for dtypes in [dh.all_int_dtypes, dh.float_dtypes]:
+            if _from in dtypes:
+                same_family = to in dtypes
+                break
+        if same_family:
+            from_min, from_max = dh.dtype_ranges[_from]
+            to_min, to_max = dh.dtype_ranges[to]
+            expected = from_min >= to_min and from_max <= to_max
+        else:
+            expected = False
+    assert out == expected, f"{out=}, but should be {expected} {f_func}"
+
 
 def make_dtype_id(dtype: DataType) -> str:
     return dh.dtype_to_name[dtype]