Expand a multidimentional array with another array of different shape

January 10, 2023

I have the following arrays:

A = np.array([
            [[[0, 1, 2, 3],
              [3, 0, 1, 2],
              [2, 3, 0, 1],
              [1, 3, 2, 1],
              [1, 2, 3, 0]]],
            
            [[[9, 8, 7, 6],
              [5, 4, 3, 2],
              [0, 9, 8, 3],
              [1, 9, 2, 3],
              [1, 0, -1, 2]]],
            
            [[[0, 7, 1, 2],
              [1, 2, 1, 0],
              [0, 2, 0, 7],
              [-1, 3, 0, 1],
              [1, 0, 1, 0]]]
              ])

A.shape
(3,1,5,4)

B = np.array([
    [[[1, 0],
      [-1, 2],
      [9, 1],
      [8, 2],
      [7, 0]]],
            
      [[[9, 6],
       [5, 2],
       [0, 3],
       [1, 9],
       [1, 0]]],
            
      [[[0, 7],
        [1, 0],
        [0, 7],
        [-1, 1],
        [0, 0]]]
])

B.shape
(3,1,5,2)

Then I want to expand array A with B in the last dimension of A. Such that, the result X is:

X = np.array([
            [[[0, 1, 2, 3, 1, 0],
              [3, 0, 1, 2,-1, 2],
              [2, 3, 0, 1, 9, 1],
              [1, 3, 2, 1, 8, 2],
              [1, 2, 3, 0, 7, 0]]],
            
            [[[9, 8, 7, 6, 9, 6],
              [5, 4, 3, 2, 5, 2],
              [0, 9, 8, 3, 0, 3],
              [1, 9, 2, 3, 1, 9],
              [1, 0,-1, 2, 1, 0]]],
            
            [[[0, 7, 1, 2, 0, 7],
              [1, 2, 1, 0, 1, 0],
              [0, 2, 0, 7, 0, 7],
              [-1,3, 0, 1,-1, 1],
              [1, 0, 1, 0, 0, 0]]]
              ])

X.shape
(3,1,5,6)
``