WIP

Author: ev-br

array_api_tests/test_manipulation_functions.py

@@ -167,7 +167,34 @@ def test_expand_dims(x, axis):
     )
 )
 def test_expand_dims_tuples(x, axes):
-    print(x.shape, axes)
+    x_ndim, y_ndim = x.ndim, x.ndim + len(axes)
+
+    # normalize
+    n_axes = tuple(ax + y_ndim if ax < 0 else ax for ax in axes)
+    unique_axes = set(n_axes)
+
+#    print(x.shape, axes, n_axes)
+
+    if any(ax < 0 or ax >= y_ndim for ax in n_axes) or len(n_axes) != len(unique_axes):
+#        print("\t raises")
+        with pytest.raises((IndexError, ValueError)):
+            xp.expand_dims(x, axis=axes)
+        return
+
+
+    repro_snippet = ph.format_snippet(f"xp.expand_dims({x!r}, axis={axes!r})")
+    try:
+        y = xp.expand_dims(x, axis=axes)
+
+        ye = x
+        for ax in sorted(n_axes):
+            ye = xp.expand_dims(ye, axis=ax)
+
+        assert y.shape == ye.shape
+
+    except Exception as exc:
+        ph.add_note(exc, repro_snippet)
+        raise
 
 
 @pytest.mark.min_version("2023.12")