numpy · eliegoudout · Nov 12, 2023 · Nov 12, 2023 · Nov 12, 2023 · Nov 17, 2023
diff --git a/doc/release/upcoming_changes/25170.new-feature.rst b/doc/release/upcoming_changes/25170.new-feature.rst
@@ -0,0 +1,7 @@
+``fill`` option for `numpy.atleast_2d` and `numpy.atleast_3d`
+-------------------------------------------------------------
+The ``fill`` option is now available for `numpy.atleast_2d` and
+`numpy.atleast_3d`, to specify where the lacking dimensions should be added.
+Previous default behaviours remain unchanged, they respectively correspond to
+``fill='left'`` for `numpy.atleast_2d` and ``fill='both'`` for
+`numpy.atleast_3d`.
diff --git a/numpy/_core/shape_base.py b/numpy/_core/shape_base.py
@@ -74,12 +74,12 @@ def atleast_1d(*arys):
         return tuple(res)
 
 
-def _atleast_2d_dispatcher(*arys):
+def _atleast_2d_dispatcher(*arys, fill=None):
     return arys
 
 
 @array_function_dispatch(_atleast_2d_dispatcher)
-def atleast_2d(*arys):
+def atleast_2d(*arys, fill='left'):
     """
     View inputs as arrays with at least two dimensions.
 
@@ -89,6 +89,8 @@ def atleast_2d(*arys):
         One or more array-like sequences.  Non-array inputs are converted
         to arrays.  Arrays that already have two or more dimensions are
         preserved.
+    fill : {'left', 'right'}
+        Where to add the missing dimension(s) to the original shape.
 
     Returns
     -------
@@ -107,8 +109,10 @@ def atleast_2d(*arys):
     array([[3.]])
 
     >>> x = np.arange(3.0)
-    >>> np.atleast_2d(x)
-    array([[0., 1., 2.]])
+    >>> np.atleast_2d(x, fill='right')
+    array([[ 0.],
+           [ 1.],
+           [ 2.]])
     >>> np.atleast_2d(x).base is x
     True
 
@@ -122,7 +126,8 @@ def atleast_2d(*arys):
         if ary.ndim == 0:
             result = ary.reshape(1, 1)
         elif ary.ndim == 1:
-            result = ary[_nx.newaxis, :]
+            fill_right = fill == 'right'
+            result = ary[:, _nx.newaxis] if fill_right else ary[_nx.newaxis]
         else:
             result = ary
         res.append(result)
@@ -132,12 +137,12 @@ def atleast_2d(*arys):
         return tuple(res)
 
 
-def _atleast_3d_dispatcher(*arys):
+def _atleast_3d_dispatcher(*arys, fill=None):
     return arys
 
 
 @array_function_dispatch(_atleast_3d_dispatcher)
-def atleast_3d(*arys):
+def atleast_3d(*arys, fill='both'):
     """
     View inputs as arrays with at least three dimensions.
 
@@ -147,6 +152,10 @@ def atleast_3d(*arys):
         One or more array-like sequences.  Non-array inputs are converted to
         arrays.  Arrays that already have three or more dimensions are
         preserved.
+    fill : {'both', 'left', 'right'}
+        Where to add the missing dimension(s) to the original shape. When
+        ``'both'``, shape ``(N,)`` gives shape ``(1, N, 1)`` and shape
+        ``(M, N)`` gives shape ``(M, N, 1)``.
 
     Returns
     -------
@@ -171,8 +180,8 @@ def atleast_3d(*arys):
     (1, 3, 1)
 
     >>> x = np.arange(12.0).reshape(4,3)
-    >>> np.atleast_3d(x).shape
-    (4, 3, 1)
+    >>> np.atleast_3d(x, fill='left').shape
+    (1, 4, 3)
     >>> np.atleast_3d(x).base is x.base  # x is a reshape, so not base itself
     True
 
@@ -192,9 +201,15 @@ def atleast_3d(*arys):
         if ary.ndim == 0:
             result = ary.reshape(1, 1, 1)
         elif ary.ndim == 1:
-            result = ary[_nx.newaxis, :, _nx.newaxis]
+            if fill == 'left':
+                result = ary[_nx.newaxis, _nx.newaxis]
+            elif fill == 'right':
+                result = ary[:, _nx.newaxis, _nx.newaxis]
+            else:
+                result = ary[_nx.newaxis, :, _nx.newaxis]
         elif ary.ndim == 2:
-            result = ary[:, :, _nx.newaxis]
+            fill_left = fill == 'left'
+            result = ary[_nx.newaxis] if fill_left else ary[..., _nx.newaxis]
         else:
             result = ary
         res.append(result)

diff --git a/numpy/_core/shape_base.pyi b/numpy/_core/shape_base.pyi
@@ -23,18 +23,26 @@ def atleast_1d(arys: ArrayLike, /) -> NDArray[Any]: ...
 def atleast_1d(*arys: ArrayLike) -> tuple[NDArray[Any], ...]: ...
 
 @overload
-def atleast_2d(arys: _ArrayLike[_SCT], /) -> NDArray[_SCT]: ...
+def atleast_2d(
+    arys: _ArrayLike[_SCT],
+    /,
+    fill: str = ...
+) -> NDArray[_SCT]: ...
 @overload
-def atleast_2d(arys: ArrayLike, /) -> NDArray[Any]: ...
+def atleast_2d(arys: ArrayLike, /, fill: str = ...) -> NDArray[Any]: ...
 @overload
-def atleast_2d(*arys: ArrayLike) -> tuple[NDArray[Any], ...]: ...
+def atleast_2d(*arys: ArrayLike, fill: str = ...) -> tuple[NDArray[Any], ...]: ...
 
 @overload
-def atleast_3d(arys: _ArrayLike[_SCT], /) -> NDArray[_SCT]: ...
+def atleast_3d(
+    arys: _ArrayLike[_SCT],
+    /,
+    fill: str = ...
+) -> NDArray[_SCT]: ...
 @overload
-def atleast_3d(arys: ArrayLike, /) -> NDArray[Any]: ...
+def atleast_3d(arys: ArrayLike, /, fill: str = ...) -> NDArray[Any]: ...
 @overload
-def atleast_3d(*arys: ArrayLike) -> tuple[NDArray[Any], ...]: ...
+def atleast_3d(*arys: ArrayLike, fill: str = ...) -> tuple[NDArray[Any], ...]: ...
 
 @overload
 def vstack(