Numpy common slicing for different dimensions

I have a function performs an operation on 2D arrays. I would like to expand its functionality to be able to work with 3D arrays, for which I would like to know if there is a way to access the vectors of each 2D array regardless of whether the input is a 3D array or a 2D array.

For example, if I have the following 2D matrix:

>>>arr2d=array([[0, 0],
                [1, 1]])

I can access the last vector using:

>>>arr2d[-1]
array([1, 1])

And if I have a 3D array like this:

>>>arr3d=array([[[ 0,  0],
                 [ 1,  1]],

                [[ 3,  0],
                 [ 2,  7]],

                [[ 9,  5],
                 [ 8,  6]],

                [[20,  4],
                 [ 6,  5]]])

I can access the last vector of each 2D submatrix using:

>>>arr3d[:,-1]
array([[1, 1],
       [2, 7],
       [8, 6],
       [6, 5]])

I would like to know if there is a common slice that I can use on both arrays to get the above results, i.e. something like the following, with some_slice being the same in each case:

>>>arr2d[some_slice]
array([1, 1])

>>>arr3d[some_slice]
array([[1, 1],
       [2, 7],
       [8, 6],
       [6, 5]])

>Solution :

Use Ellipsis as below:

print(arr2d[..., -1, :])
print(arr3d[..., -1, :])

Output

[1 1]
[[1 1]
 [2 7]
 [8 6]
 [6 5]]

From the documentation (emphasis mine):

Ellipsis expands to the number of : objects needed for the selection
tuple to index all dimensions. In most cases, this means that the
length of the expanded selection tuple is x.ndim. There may only be a
single ellipsis present

But if you want to create a function that works both for 2d and 3d arrays I suggest that you convert the 2d array two a 3d array by adding a new axis. Find a toy example below:

def foo_index(arr):
    if len(arr.shape) == 2:
        arr = arr[np.newaxis, :]
    return arr[:, -1]


print(foo_index(arr2d))  # two-dimensional shape
print(foo_index(arr3d))  # two-dimensional shape

Note that the output now have the same shape (2d), therefore code depending on the result can work regardless of a 2d or 3d input array. Note that this will not happen by using the same slice for both arrays.