WIP: add axis tuple support to torch.expand_dims

Author: ev-br

array_api_compat/torch/_aliases.py

@@ -690,9 +690,23 @@ def triu(x: Array, /, *, k: int = 0) -> Array:
     return torch.triu(x, k)
 
 # Functions that aren't in torch https://github.com/pytorch/pytorch/issues/58742
-def expand_dims(x: Array, /, *, axis: int = 0) -> Array:
-    return torch.unsqueeze(x, axis)
+def expand_dims(x: Array, /, *, axis: int | tuple[int, ...]) -> Array:
+    if isinstance(axis, int):
+        return torch.unsqueeze(x, axis)
+    else:
+        y_ndim = x.ndim + len(axis)
+
+        # normalize
+        n_axis = tuple(ax + y_ndim if ax < 0 else ax for ax in axis)
+        if (len(n_axis) != len(set(n_axis)) or
+            _builtin_any(ax < 0 or ax >= y_ndim for ax in n_axis)
+        ):
+            raise ValueError()
+
+        shape_it = iter(x.shape)
+        shape = [1 if ax in n_axis else next(shape_it) for ax in range(y_ndim)]
 
+        return torch.reshape(x, shape)
 
 def astype(
     x: Array,