Best way to contain information with different shapes in a numpy array?

Question

I want to implement conditional variational autoencoder for MNIST images with convolutional layers, not flattening. So the image itself is [batch_size, 28, 28, 1] (