Question: I have a (5, 5) 2D Numpy array: map_height = 5 map_width = 5 # Define a 2D np array- a = np.arange(map_height * map_width).reshape(map_height, map_width) # a ''' array([[ 0, 1, 2, 3, 4], [ 5, 6, 7, 8, 9], [10, 11, 12, 13, 14], [15, 16, 17, 18, 19], [20, 21, 22, 23, 24]]) '''  I can wrap this array around on both of its axes using 'pad()': a_wrapped = np.pad(array = a, pad_width = 1, mode = 'wrap') a_wrapped ''' array([[24, 20, 21, 22, 23, 24, 20], [ 4, 0, 1, 2, 3, 4, 0], [ 9, 5, 6, 7, 8, 9, 5], [14, 10, 11, 12, 13, 14, 10], [19, 15, 16, 17, 18, 19, 15], [24, 20, 21, 22, 23, 24, 20], [ 4, 0, 1, 2, 3, 4, 0]]) '''  The 2D coordinates of (5, 5) 'a' are computed (inefficiently) as: # 2D coordinates - # 1st channel/axis = row indices & 2 channel/axis = column indices. a_2d_coords = np.zeros((map_height, map_width, 2), dtype = np.int16) for row_idx in range(map_height): for col_idx in range(map_width): a_2d_coords[row_idx, col_idx][0] = row_idx a_2d_coords[row_idx, col_idx][1] = col_idx # a_2d_coords.shape # (5, 5, 2)  I want to wrap this 2D coordinates array 'a_2d_coords' as well by doing: a_2d_coords_wrapped = np.pad(array = a_2d_coords, pad_width = 1, mode = 'wrap') # a_2d_coords_wrapped.shape # (7, 7, 4)  It also wraps the 3rd axis/dimension which should not be done! The goal is that the coords of a[1, 4] = (1, 4) and its neighbor's coords to right hand side (RHS) should be a[1, 0] = (1, 0). This is wrapping around the x-axis. Similarly, y-axis 2D coordinates should also be wrapped.

Answer: First of all, you can simplify your calculation of the 2D coordinates as follows: a_2d_coords = np.moveaxis(np.mgrid[:map_height, :map_width], 0, -1)  Let's check if this still produces the same result: import numpy as np # For the check, use different height and width values to detect swaps map_height, map_width = 3, 7 # Original implementation for coords a_2d_coords_given = np.zeros((map_height, map_width, 2), dtype = np.int16) for row_idx in range(map_height): for col_idx in range(map_width): a_2d_coords_given[row_idx, col_idx][0] = row_idx a_2d_coords_given[row_idx, col_idx][1] = col_idx # Proposed implementation for coords a_2d_coords = np.moveaxis(np.mgrid[:map_height, :map_width], 0, -1) # Equality check assert np.all(a_2d_coords == a_2d_coords_given)  Second, if you note that np.pad's pad_width argument can also take axis-specific values, then you can achieve the desired coordinate padding as follows: a_2d_coords_wrapped = np.pad(a_2d_coords, mode="wrap", pad_width=((1, 1), (1, 1), (0, 0)))  Now we pad both axis 0 and axis 1 with one leading and one trailing value ((1, 1)), but leave axis 2 alone ((0, 0)). Altogether, we have: import numpy as np map_height, map_width = 5, 5 a = np.arange(map_height * map_width).reshape(map_height, map_width) a_2d_coords = np.moveaxis(np.mgrid[:map_height, :map_width], 0, -1) a_2d_coords_wrapped = np.pad(a_2d_coords, mode="wrap", pad_width=((1, 1), (1, 1), (0, 0))) print(a[a_2d_coords_wrapped[..., 0], a_2d_coords_wrapped[..., 1]]) # Prints # [[24 20 21 22 23 24 20] # [ 4 0 1 2 3 4 0] # [ 9 5 6 7 8 9 5] # [14 10 11 12 13 14 10] # [19 15 16 17 18 19 15] # [24 20 21 22 23 24 20] # [ 4 0 1 2 3 4 0]]  Further simplification We can simplify the code even one step further by providing a_2d_coords/a_2d_coords_wrapped as a 2×H×W array rather than a H×W×2 array: import numpy as np map_height, map_width = 5, 5 a = np.arange(map_height * map_width).reshape(map_height, map_width) a_2d_coords = np.mgrid[:map_height, :map_width] a_2d_coords_wrapped = np.pad(a_2d_coords, mode="wrap", pad_width=((0, 0), (1, 1), (1, 1))) print(a[*a_2d_coords_wrapped]) # Prints # [[24 20 21 22 23 24 20] # [ 4 0 1 2 3 4 0] # [ 9 5 6 7 8 9 5] # [14 10 11 12 13 14 10] # [19 15 16 17 18 19 15] # [24 20 21 22 23 24 20] # [ 4 0 1 2 3 4 0]]  Mind the * in a[*a_2d_coords_wrapped], which is necessary for "unpacking" the coordinates, so that two H×W arrays rather than one 2×H×W array are provided for indexing.