Smoke all manipulation methods

Author: honno

array_api_tests/test_manipulation_functions.py

@@ -9,6 +9,8 @@
 from . import pytest_helpers as ph
 from . import xps
 
+shared_shapes = st.shared(hh.shapes(), key="shape")
+
 
 @given(
     shape=hh.shapes(min_dims=1),
@@ -32,6 +34,81 @@ def test_concat(shape, dtypes, kw, data):
     # TODO: assert out elements match input arrays
 
 
+@given(
+    x=xps.arrays(dtype=xps.scalar_dtypes(), shape=shared_shapes),
+    axis=shared_shapes.flatmap(lambda s: st.integers(-len(s), len(s))),
+)
+def test_expand_dims(x, axis):
+    xp.expand_dims(x, axis=axis)
+    # TODO
+
+
+@given(
+    x=xps.arrays(dtype=xps.scalar_dtypes(), shape=shared_shapes),
+    kw=hh.kwargs(
+        axis=st.one_of(
+            st.none(),
+            shared_shapes.flatmap(
+                lambda s: st.none()
+                if len(s) == 0
+                else st.integers(-len(s) + 1, len(s) - 1),
+            ),
+        )
+    ),
+)
+def test_flip(x, kw):
+    xp.flip(x, **kw)
+    # TODO
+
+
+@given(
+    x=xps.arrays(dtype=xps.scalar_dtypes(), shape=shared_shapes),
+    axes=shared_shapes.flatmap(
+        lambda s: st.lists(
+            st.integers(0, max(len(s) - 1, 0)),
+            min_size=len(s),
+            max_size=len(s),
+            unique=True,
+        ).map(tuple)
+    ),
+)
+def test_permute_dims(x, axes):
+    xp.permute_dims(x, axes)
+    # TODO
+
+
+@given(
+    x=xps.arrays(dtype=xps.scalar_dtypes(), shape=shared_shapes),
+    shape=shared_shapes,  # TODO: test more compatible shapes
+)
+def test_reshape(x, shape):
+    xp.reshape(x, shape)
+    # TODO
+
+
+@given(
+    # TODO: axis arguments, update shift respectively
+    x=xps.arrays(dtype=xps.scalar_dtypes(), shape=shared_shapes),
+    shift=shared_shapes.flatmap(lambda s: st.integers(0, max(math.prod(s) - 1, 0))),
+)
+def test_roll(x, shift):
+    xp.roll(x, shift)
+    # TODO
+
+
+@given(
+    x=xps.arrays(dtype=xps.scalar_dtypes(), shape=shared_shapes),
+    axis=shared_shapes.flatmap(
+        lambda s: st.just(0)
+        if len(s) == 0
+        else st.integers(-len(s) + 1, len(s) - 1).filter(lambda i: s[i] == 1)
+    ),  # TODO: tuple of axis i.e. axes
+)
+def test_squeeze(x, axis):
+    xp.squeeze(x, axis)
+    # TODO
+
+
 @given(
     shape=hh.shapes(),
     dtypes=hh.mutually_promotable_dtypes(None),