Numpy: stack arrays whose internal dimensions differ

To follow up on Divakar's answer above, the axis argument in numpy is the index of a given dimension within an array's shape. Here I want to stack a and b by virtue of their middle shape value, which is at index = 1:

import numpy as np

a = np.random.rand(5, 1, 3)
b = np.random.rand(5, 100, 3)

# create the desired result shape: 55, 627, 3
stacked = np.concatenate((b, a), axis=1)

# validate that a was appended to the end of b
print(stacked[:, -1, :], '\n\n\n', a.squeeze())

This returns:

[[0.72598529 0.99395887 0.21811998]
 [0.9833895  0.465955   0.29518207]
 [0.38914048 0.61633291 0.0132326 ]
 [0.05986115 0.81354865 0.43589306]
 [0.17706517 0.94801426 0.4567973 ]] 


 [[0.72598529 0.99395887 0.21811998]
 [0.9833895  0.465955   0.29518207]
 [0.38914048 0.61633291 0.0132326 ]
 [0.05986115 0.81354865 0.43589306]
 [0.17706517 0.94801426 0.4567973 ]]

A purist might use instead np.all(stacked[:, -1, :] == a.squeeze()) to validate this equivalence. All glory to @Divakar!

Strictly for the curious, the use case for this concatenation is a kind of wonky data preparation pipeline for a Long Short Term Memory Neural Network. In that kind of network, the training data shape should be number_of_observations, number_of_time_intervals, number_of_dimensions_per_observation. I am generating new predictions of each object at a new time interval, so those predictions have shape number_of_observations, 1, number_of_dimensions_per_observation. To visualize the sequence of observations' positions over time, I want to add the new positions to the array of previous positions, hence the question above.