From 2474c4717a40f37eaf5e5a06f608403c5d5c6672 Mon Sep 17 00:00:00 2001
From: "Joseph R. Fox-Rabinovitz" <jfoxrabinovitz@gmail.com>
Date: Wed, 10 Feb 2021 05:21:35 -0500
Subject: [PATCH] ENH: Added atleast_nd to numpy and ma

---
 .../upcoming_changes/18386.new_function.rst   |   4 +
 .../reference/routines.array-manipulation.rst |   1 +
 doc/source/reference/routines.ma.rst          |   1 +
 doc/source/user/quickstart.rst                |   1 +
 numpy/__init__.pyi                            |   1 +
 numpy/core/shape_base.py                      | 190 ++++++++++++------
 numpy/core/shape_base.pyi                     |   2 +
 numpy/core/tests/test_shape_base.py           |  71 ++++++-
 numpy/lib/shape_base.py                       |   2 +-
 numpy/ma/__init__.pyi                         |   1 +
 numpy/ma/extras.py                            |   9 +-
 numpy/ma/tests/test_extras.py                 |  44 +++-
 .../tests/data/pass/array_constructors.py     |   3 +
 .../tests/data/reveal/array_constructors.py   |   5 +
 14 files changed, 262 insertions(+), 73 deletions(-)
 create mode 100644 doc/release/upcoming_changes/18386.new_function.rst

diff --git a/doc/release/upcoming_changes/18386.new_function.rst b/doc/release/upcoming_changes/18386.new_function.rst
new file mode 100644
index 000000000000..33925a931f31
--- /dev/null
+++ b/doc/release/upcoming_changes/18386.new_function.rst
@@ -0,0 +1,4 @@
+Added `atleast_nd` function to `numpy` and `numpy.ma`
+-----------------------------------------------------
+`atleast_nd` generalizes ``atleast_1d``, ``atleast_2d`` and ``atleast_3d`` to
+arbitrary numbers of dimensions.
diff --git a/doc/source/reference/routines.array-manipulation.rst b/doc/source/reference/routines.array-manipulation.rst
index 1c96495d96f7..c48346757a96 100644
--- a/doc/source/reference/routines.array-manipulation.rst
+++ b/doc/source/reference/routines.array-manipulation.rst
@@ -41,6 +41,7 @@ Changing number of dimensions
    atleast_1d
    atleast_2d
    atleast_3d
+   atleast_nd
    broadcast
    broadcast_to
    broadcast_arrays
diff --git a/doc/source/reference/routines.ma.rst b/doc/source/reference/routines.ma.rst
index d961cbf02f6c..bbd663aba444 100644
--- a/doc/source/reference/routines.ma.rst
+++ b/doc/source/reference/routines.ma.rst
@@ -127,6 +127,7 @@ Changing the number of dimensions
    ma.atleast_1d
    ma.atleast_2d
    ma.atleast_3d
+   ma.atleast_nd
    ma.expand_dims
    ma.squeeze
 
diff --git a/doc/source/user/quickstart.rst b/doc/source/user/quickstart.rst
index ab5bb5318cdd..8a54044d1c56 100644
--- a/doc/source/user/quickstart.rst
+++ b/doc/source/user/quickstart.rst
@@ -967,6 +967,7 @@ Conversions
    `atleast_1d`,
    `atleast_2d`,
    `atleast_3d`,
+   `atleast_nd`,
    `mat`
 Manipulations
     `array_split`,
diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi
index 351cacffaa36..875fb847f64e 100644
--- a/numpy/__init__.pyi
+++ b/numpy/__init__.pyi
@@ -326,6 +326,7 @@ from numpy.core.shape_base import (
     atleast_1d as atleast_1d,
     atleast_2d as atleast_2d,
     atleast_3d as atleast_3d,
+    atleast_nd as atleast_nd,
     block as block,
     hstack as hstack,
     stack as stack,
diff --git a/numpy/core/shape_base.py b/numpy/core/shape_base.py
index 89e98ab3072d..2584bb5b8eef 100644
--- a/numpy/core/shape_base.py
+++ b/numpy/core/shape_base.py
@@ -1,5 +1,5 @@
-__all__ = ['atleast_1d', 'atleast_2d', 'atleast_3d', 'block', 'hstack',
-           'stack', 'vstack']
+__all__ = ['atleast_1d', 'atleast_2d', 'atleast_3d', 'atleast_nd',
+           'block', 'hstack', 'stack', 'vstack']
 
 import functools
 import itertools
@@ -17,6 +17,12 @@
     overrides.array_function_dispatch, module='numpy')
 
 
+def _unpack(lst):
+    if len(lst) == 1:
+        return lst[0]
+    return lst
+
+
 def _atleast_1d_dispatcher(*arys):
     return arys
 
@@ -42,7 +48,7 @@ def atleast_1d(*arys):
 
     See Also
     --------
-    atleast_2d, atleast_3d
+    atleast_2d, atleast_3d, atleast_nd
 
     Examples
     --------
@@ -61,18 +67,7 @@ def atleast_1d(*arys):
     [array([1]), array([3, 4])]
 
     """
-    res = []
-    for ary in arys:
-        ary = asanyarray(ary)
-        if ary.ndim == 0:
-            result = ary.reshape(1)
-        else:
-            result = ary
-        res.append(result)
-    if len(res) == 1:
-        return res[0]
-    else:
-        return res
+    return _unpack([atleast_nd(a, 1, 0) for a in arys])
 
 
 def _atleast_2d_dispatcher(*arys):
@@ -87,9 +82,9 @@ def atleast_2d(*arys):
     Parameters
     ----------
     arys1, arys2, ... : array_like
-        One or more array-like sequences.  Non-array inputs are converted
-        to arrays.  Arrays that already have two or more dimensions are
-        preserved.
+        One or more array-like sequences.  Non-array inputs are
+        converted to arrays. Arrays that already have two or more
+        dimensions are preserved.
 
     Returns
     -------
@@ -100,7 +95,7 @@ def atleast_2d(*arys):
 
     See Also
     --------
-    atleast_1d, atleast_3d
+    atleast_1d, atleast_3d, atleast_nd
 
     Examples
     --------
@@ -117,20 +112,7 @@ def atleast_2d(*arys):
     [array([[1]]), array([[1, 2]]), array([[1, 2]])]
 
     """
-    res = []
-    for ary in arys:
-        ary = asanyarray(ary)
-        if ary.ndim == 0:
-            result = ary.reshape(1, 1)
-        elif ary.ndim == 1:
-            result = ary[_nx.newaxis, :]
-        else:
-            result = ary
-        res.append(result)
-    if len(res) == 1:
-        return res[0]
-    else:
-        return res
+    return _unpack([atleast_nd(a, 2, 0) for a in arys])
 
 
 def _atleast_3d_dispatcher(*arys):
@@ -145,22 +127,31 @@ def atleast_3d(*arys):
     Parameters
     ----------
     arys1, arys2, ... : array_like
-        One or more array-like sequences.  Non-array inputs are converted to
-        arrays.  Arrays that already have three or more dimensions are
-        preserved.
+        One or more array-like sequences.  Non-array inputs are
+        converted to arrays.  Arrays that already have three or more
+        dimensions are preserved.
 
     Returns
     -------
     res1, res2, ... : ndarray
-        An array, or list of arrays, each with ``a.ndim >= 3``.  Copies are
-        avoided where possible, and views with three or more dimensions are
-        returned.  For example, a 1-D array of shape ``(N,)`` becomes a view
-        of shape ``(1, N, 1)``, and a 2-D array of shape ``(M, N)`` becomes a
-        view of shape ``(M, N, 1)``.
+        An array, or list of arrays, each with ``a.ndim >= 3``.  Copies
+        are avoided where possible, and views with three or more
+        dimensions are returned.  For example, a 1-D array of shape
+        ``(N,)`` becomes a view of shape ``(1, N, 1)``, and a 2-D array
+        of shape ``(M, N)`` becomes a view of shape ``(M, N, 1)``.
 
     See Also
     --------
-    atleast_1d, atleast_2d
+    atleast_1d, atleast_2d, atleast_nd
+
+    Notes
+    -----
+    As mentioned in the `Returns` section, the results of this
+    function are not consistent with any of the other `atleast*`
+    functions. `atleast_2d` prepends the unit dimension to a 1D array
+    while `atleast_3d` appends it to a 2D array. The 1D array case
+    both appends and prepends a dimension, while `atleast_nd` can only
+    add dimensions to one end at a time.
 
     Examples
     --------
@@ -187,22 +178,105 @@ def atleast_3d(*arys):
     [[[1 2]]] (1, 1, 2)
 
     """
-    res = []
-    for ary in arys:
-        ary = asanyarray(ary)
-        if ary.ndim == 0:
-            result = ary.reshape(1, 1, 1)
-        elif ary.ndim == 1:
-            result = ary[_nx.newaxis, :, _nx.newaxis]
-        elif ary.ndim == 2:
-            result = ary[:, :, _nx.newaxis]
-        else:
-            result = ary
-        res.append(result)
-    if len(res) == 1:
-        return res[0]
-    else:
-        return res
+    return _unpack([atleast_nd(atleast_nd(a, 2, 0), 3, -1) for a in arys])
+
+
+def _atleast_nd_dispatcher(ary, ndim, pos=None):
+    return (ary,)
+
+
+@array_function_dispatch(_atleast_nd_dispatcher)
+def atleast_nd(ary, ndim, pos=0):
+    """
+    View input as array with at least `ndim` dimensions.
+
+    New unit dimensions are inserted at the index given by `pos` if
+    necessary.
+
+    Parameters
+    ----------
+    ary : array_like
+        The input array. Non-array inputs are converted to arrays.
+        Arrays that already have `ndim` or more dimensions are
+        preserved.
+    ndim : int
+        The minimum number of dimensions required.
+    pos : int, optional
+        The index to insert the new dimensions. May range from
+        ``-ary.ndim - 1`` to ``+ary.ndim`` (inclusive). Non-negative
+        indices indicate locations before the corresponding axis:
+        ``pos=0`` means to insert at the very beginning. Negative
+        indices indicate locations after the corresponding axis:
+        ``pos=-1`` means to insert at the very end. 0 and -1 are always
+        guaranteed to work. Any other number will depend on the
+        dimensions of the existing array. Default is 0.
+
+    Returns
+    -------
+    res : ndarray
+        An array with ``res.ndim >= ndim``. A view is returned for array
+        inputs. Dimensions are prepended if `pos` is 0, so for example,
+        a 1-D array of shape ``(N,)`` with ``ndim=4`` becomes a view of
+        shape ``(1, 1, 1, N)``. Dimensions are appended if `pos` is -1,
+        so for example a 2-D array of shape ``(M, N)`` becomes a view of
+        shape ``(M, N, 1, 1)`` when ``ndim=4``.
+
+    See Also
+    --------
+    atleast_1d, atleast_2d, atleast_3d
+
+    Notes
+    -----
+    This function does not follow the convention of the other
+    ``atleast_*d`` functions in numpy in that it only accepts a single
+    array argument. To process multiple arrays, use a comprehension or
+    loop around the function call. See examples below.
+
+    Setting ``pos=0`` is equivalent to how the array would be
+    interpreted by numpy's broadcasting rules. There is no need to call
+    this function for simple broadcasting. This is also roughly
+    (but not exactly) equivalent to
+    ``np.array(ary, copy=False, subok=True, ndmin=ndim)``.
+
+    It is easy to create functions for specific dimensions similar to
+    the other ``atleast_*d`` functions using Python's
+    `functools.partial` function. An example is shown below.
+
+    Examples
+    --------
+    >>> np.atleast_nd(3.0, 4)
+    array([[[[ 3.]]]])
+
+    >>> x = np.arange(3.0)
+    >>> np.atleast_nd(x, 2).shape
+    (1, 3)
+
+    >>> x = np.arange(12.0).reshape(4, 3)
+    >>> np.atleast_nd(x, 5).shape
+    (1, 1, 1, 4, 3)
+    >>> np.atleast_nd(x, 5).base is x.base
+    True
+
+    >>> [np.atleast_nd(x, 2) for x in ((1, 2), [[3, 4]], [[[5, 6]]])]
+    [array([[1, 2]]), array([[3, 4]]), array([[[5, 6]]])]
+
+    >>> np.atleast_nd((1, 2), 5, pos=0).shape
+    (1, 1, 1, 1, 2)
+    >>> np.atleast_nd((1, 2), 5, pos=-1).shape
+    (2, 1, 1, 1, 1)
+
+    >>> from functools import partial
+    >>> atleast_4d = partial(np.atleast_nd, ndim=4)
+    >>> atleast_4d([1, 2, 3])
+    [[[[1, 2, 3]]]]
+    """
+    ary = array(ary, copy=False, subok=True)
+    pos = normalize_axis_index(pos, ary.ndim + 1)
+    extra = operator.index(ndim) - ary.ndim
+    if extra > 0:
+        ind = pos * (slice(None),) + extra * (None,) + (Ellipsis,)
+        ary = ary[ind]
+    return ary
 
 
 def _arrays_for_stack_dispatcher(arrays, stacklevel=4):
diff --git a/numpy/core/shape_base.pyi b/numpy/core/shape_base.pyi
index ec40a88143b1..e9248407475b 100644
--- a/numpy/core/shape_base.pyi
+++ b/numpy/core/shape_base.pyi
@@ -26,6 +26,8 @@ def atleast_3d(__arys: ArrayLike) -> ndarray: ...
 @overload
 def atleast_3d(*arys: ArrayLike) -> List[ndarray]: ...
 
+def atleast_nd(ary: ArrayLike, ndim: int, pos: int = ...) -> ndarray: ...
+
 def vstack(tup: Sequence[ArrayLike]) -> ndarray: ...
 def hstack(tup: Sequence[ArrayLike]) -> ndarray: ...
 @overload
diff --git a/numpy/core/tests/test_shape_base.py b/numpy/core/tests/test_shape_base.py
index a0c72f9d0fcb..bc3033ab6b0b 100644
--- a/numpy/core/tests/test_shape_base.py
+++ b/numpy/core/tests/test_shape_base.py
@@ -1,9 +1,9 @@
+from decimal import Decimal
 import pytest
 import numpy as np
 from numpy.core import (
-    array, arange, atleast_1d, atleast_2d, atleast_3d, block, vstack, hstack,
-    newaxis, concatenate, stack
-    )
+    array, arange, atleast_1d, atleast_2d, atleast_3d, atleast_nd,
+    block, vstack, hstack, newaxis, concatenate, stack)
 from numpy.core.shape_base import (_block_dispatcher, _block_setup,
                                    _block_concatenate, _block_slicing)
 from numpy.testing import (
@@ -123,6 +123,71 @@ def test_3D_array(self):
         assert_array_equal(res, desired)
 
 
+class TestAtleastNd(object):
+    def test_0D_arrays(self):
+        a = array(3)
+        dims = [3, 2, 0]
+        expected = [array([[[3]]]), array([[3]]), array(3)]
+
+        for b, d in zip(expected, dims):
+            assert_array_equal(atleast_nd(a, d), b)
+            assert_array_equal(atleast_nd(a, d, -1), b)
+
+    def test_nD_arrays(self):
+        a = array([1])
+        b = array([4, 5, 6])
+        c = array([[2, 3]])
+        d = array([[[2], [3]], [[2], [3]]])
+        e = ((((1, 2), (3, 4)), ((5, 6), (7, 8))))
+        arrays = (a, b, c, d, e)
+        expected_before = (array([[[1]]]),
+                           array([[[4, 5, 6]]]),
+                           array([[[2, 3]]]),
+                           d,
+                           array(e))
+        expected_after = (array([[[1]]]),
+                          array([[[4]], [[5]], [[6]]]),
+                          array([[[2], [3]]]),
+                          d,
+                          array(e))
+
+        for x, y in zip(arrays, expected_before):
+            assert_array_equal(atleast_nd(x, 3), y)
+        for x, y in zip(arrays, expected_after):
+            assert_array_equal(atleast_nd(x, 3, pos=-1), y)
+
+    def test_nocopy(self):
+        a = arange(12.0).reshape(4, 3)
+        res = atleast_nd(a, 5)
+        desired_shape = (1, 1, 1, 4, 3)
+        desired_base = a.base  # a was reshaped
+        assert_equal(res.shape, desired_shape)
+        assert_(res.base is desired_base)
+
+    def test_passthough(self):
+        a = array([1, 2, 3])
+        assert_(atleast_nd(a, 0) is a)
+        assert_(atleast_nd(a, 1) is a)
+
+    def test_other_pos(self):
+        a = arange(12.0).reshape(4, 3)
+        res = atleast_nd(a, 4, pos=1)
+        assert_equal(res.shape, (4, 1, 1, 3))
+        assert_raises(ValueError, atleast_nd, a, 4, pos=5)
+
+    def test_ndim(self):
+        a = 3
+        assert_raises(TypeError, atleast_nd, a, 0.4)
+        assert_raises(TypeError, atleast_nd, a, Decimal(4))
+        assert_raises(TypeError, atleast_nd, a, np.array([4, 5]))
+        assert_raises(np.AxisError, atleast_nd, a, -2, 1)
+        assert_equal(atleast_nd(a, np.array(4, dtype=np.uint8)).ndim, 4)
+
+        assert isinstance(atleast_nd(a, 0, 0), np.ndarray)
+        assert_equal(atleast_nd(a, -5).ndim, 0)
+        assert_equal(atleast_nd(a, -5, -1).ndim, 0)
+
+
 class TestHstack:
     def test_non_iterable(self):
         assert_raises(TypeError, hstack, 1)
diff --git a/numpy/lib/shape_base.py b/numpy/lib/shape_base.py
index 9dfeee527c1a..11b8bef8e7c5 100644
--- a/numpy/lib/shape_base.py
+++ b/numpy/lib/shape_base.py
@@ -542,7 +542,7 @@ def expand_dims(a, axis):
     --------
     squeeze : The inverse operation, removing singleton dimensions
     reshape : Insert, remove, and combine dimensions, and resize existing ones
-    doc.indexing, atleast_1d, atleast_2d, atleast_3d
+    doc.indexing, atleast_1d, atleast_2d, atleast_3d, atleast_nd
 
     Examples
     --------
diff --git a/numpy/ma/__init__.pyi b/numpy/ma/__init__.pyi
index 66dfe40de6a5..08c96c86cf36 100644
--- a/numpy/ma/__init__.pyi
+++ b/numpy/ma/__init__.pyi
@@ -185,6 +185,7 @@ apply_over_axes: Any
 atleast_1d: Any
 atleast_2d: Any
 atleast_3d: Any
+atleast_nd: Any
 average: Any
 clump_masked: Any
 clump_unmasked: Any
diff --git a/numpy/ma/extras.py b/numpy/ma/extras.py
index a775a15bfcfc..f6fe6d39faa4 100644
--- a/numpy/ma/extras.py
+++ b/numpy/ma/extras.py
@@ -10,7 +10,7 @@
 """
 __all__ = [
     'apply_along_axis', 'apply_over_axes', 'atleast_1d', 'atleast_2d',
-    'atleast_3d', 'average', 'clump_masked', 'clump_unmasked',
+    'atleast_3d', 'atleast_nd', 'average', 'clump_masked', 'clump_unmasked',
     'column_stack', 'compress_cols', 'compress_nd', 'compress_rowcols',
     'compress_rows', 'count_masked', 'corrcoef', 'cov', 'diagflat', 'dot',
     'dstack', 'ediff1d', 'flatnotmasked_contiguous', 'flatnotmasked_edges',
@@ -272,12 +272,10 @@ def __call__(self, x, *args, **params):
         func = getattr(np, self.__name__)
         if isinstance(x, ndarray):
             _d = func(x.__array__(), *args, **params)
-            _m = func(getmaskarray(x), *args, **params)
-            return masked_array(_d, mask=_m)
         else:
             _d = func(np.asarray(x), *args, **params)
-            _m = func(getmaskarray(x), *args, **params)
-            return masked_array(_d, mask=_m)
+        _m = func(getmaskarray(x), *args, **params)
+        return masked_array(_d, mask=_m)
 
 
 class _fromnxfunction_seq(_fromnxfunction):
@@ -344,6 +342,7 @@ def __call__(self, *args, **params):
 atleast_1d = _fromnxfunction_allargs('atleast_1d')
 atleast_2d = _fromnxfunction_allargs('atleast_2d')
 atleast_3d = _fromnxfunction_allargs('atleast_3d')
+atleast_nd = _fromnxfunction_single('atleast_nd')
 
 vstack = row_stack = _fromnxfunction_seq('vstack')
 hstack = _fromnxfunction_seq('hstack')
diff --git a/numpy/ma/tests/test_extras.py b/numpy/ma/tests/test_extras.py
index d237829cb71f..1c456c32ff16 100644
--- a/numpy/ma/tests/test_extras.py
+++ b/numpy/ma/tests/test_extras.py
@@ -23,12 +23,13 @@
     nomask, ones, zeros, count
     )
 from numpy.ma.extras import (
-    atleast_1d, atleast_2d, atleast_3d, mr_, dot, polyfit, cov, corrcoef,
-    median, average, unique, setxor1d, setdiff1d, union1d, intersect1d, in1d,
-    ediff1d, apply_over_axes, apply_along_axis, compress_nd, compress_rowcols,
-    mask_rowcols, clump_masked, clump_unmasked, flatnotmasked_contiguous,
-    notmasked_contiguous, notmasked_edges, masked_all, masked_all_like, isin,
-    diagflat, stack, vstack
+    apply_along_axis, apply_over_axes, atleast_1d, atleast_2d, atleast_3d,
+    atleast_nd, average, clump_masked, clump_unmasked, compress_nd,
+    compress_rowcols, corrcoef, cov, diagflat, dot, ediff1d,
+    flatnotmasked_contiguous, in1d, intersect1d, isin, mask_rowcols,
+    masked_all, masked_all_like, median, mr_, notmasked_contiguous,
+    notmasked_edges, polyfit, setdiff1d, setxor1d, stack, union1d, unique,
+    vstack
     )
 
 
@@ -1562,6 +1563,25 @@ def test_atleast_2d(self):
         assert_equal(a.mask.shape, a.data.shape)
         assert_equal(b.mask.shape, b.data.shape)
 
+    def test_atleast_nd(self):
+        # Test atleast_nd
+        a = masked_array([0, 1, 2], mask=[0, 1, 0])
+        b = atleast_nd(a, 4)
+        c = atleast_nd(a, 4, pos=-1)
+        d = atleast_nd(a, 4, pos=1)
+        assert_equal(a.shape, (3,))
+        assert_equal(a.data.shape, a.shape)
+        assert_equal(a.mask.shape, a.shape)
+        assert_equal(b.shape, (1, 1, 1, 3))
+        assert_equal(b.data.shape, b.shape)
+        assert_equal(b.mask.shape, b.shape)
+        assert_equal(c.shape, (3, 1, 1, 1))
+        assert_equal(c.data.shape, c.shape)
+        assert_equal(c.mask.shape, c.shape)
+        assert_equal(d.shape, (3, 1, 1, 1))
+        assert_equal(d.data.shape, d.shape)
+        assert_equal(d.mask.shape, d.shape)
+
     def test_shape_scalar(self):
         # the atleast and diagflat function should work with scalars
         # GitHub issue #3367
@@ -1600,6 +1620,18 @@ def test_shape_scalar(self):
             assert_equal(a.mask.shape, a.shape)
             assert_equal(a.data.shape, a.shape)
 
+        b = atleast_nd(1.0, 4)
+        c = atleast_nd(1.0, 4, -1)
+        d = atleast_nd(1.0, 0)
+        e = atleast_nd(1.0, -3)
+        assert_equal(b.shape, (1, 1, 1, 1))
+        assert_equal(b.mask.shape, b.data.shape)
+        assert_equal(c.shape, (1, 1, 1, 1))
+        assert_equal(c.mask.shape, c.data.shape)
+        assert_equal(d.shape, ())
+        assert_equal(d.mask.shape, d.data.shape)
+        assert_equal(e.shape, ())
+        assert_equal(e.mask.shape, e.data.shape)
 
         b = diagflat(1.0)
         assert_equal(b.shape, (1, 1))
diff --git a/numpy/typing/tests/data/pass/array_constructors.py b/numpy/typing/tests/data/pass/array_constructors.py
index 63208f139c39..264ffffa9f1c 100644
--- a/numpy/typing/tests/data/pass/array_constructors.py
+++ b/numpy/typing/tests/data/pass/array_constructors.py
@@ -114,6 +114,9 @@ def func(i: int, j: int, **kwargs: Any) -> SubClass:
 
 np.atleast_3d(C)
 
+np.atleast_nd(C, ndim=2)
+np.atleast_nd(C, ndim=2, pos=1)
+
 np.vstack([C, C])
 np.vstack([C, A])
 np.vstack([A, A])
diff --git a/numpy/typing/tests/data/reveal/array_constructors.py b/numpy/typing/tests/data/reveal/array_constructors.py
index 04d5cd229e4b..accd63bf7d46 100644
--- a/numpy/typing/tests/data/reveal/array_constructors.py
+++ b/numpy/typing/tests/data/reveal/array_constructors.py
@@ -88,6 +88,11 @@ def func(i: int, j: int, **kwargs: Any) -> SubClass: ...
 
 reveal_type(np.atleast_3d(A))  # E: numpy.ndarray[Any, Any]
 
+reveal_type(np.atleast_nd(A, 0))  # E: numpy.ndarray[Any, Any]
+reveal_type(np.atleast_nd(A, 2))  # E: numpy.ndarray[Any, Any]
+reveal_type(np.atleast_nd(B, 2))  # E: numpy.ndarray[Any, Any]
+reveal_type(np.atleast_nd(C, 2))  # E: numpy.ndarray[Any, Any]
+
 reveal_type(np.vstack([A, A]))  # E: numpy.ndarray[Any, Any]
 reveal_type(np.vstack([A, C]))  # E: numpy.ndarray[Any, Any]
 reveal_type(np.vstack([C, C]))  # E: numpy.ndarray[Any, Any]