Olá Victor!
Para passar para a rede neural as imagens são geralmente convertidas para vetores (um array unidimensional), pois essas camadas trabalham com dados em formato vetorial. Por isso que anteriormente você viu ser feito dessa forma.
No entanto, para visualização e algumas operações específicas é mais conveniente e necessário trabalhar com imagens no formato matricial (bidimensional). No código que você mencionou, o X_teste[0].reshape(28,28) está convertendo o vetor de volta para uma matriz de 28×28 pixels para que possamos visualizar a imagem original de maneira mais intuitiva. Ou seja, é apenas para visualizarmos a imagem (repare que para isso estamos usando a função plt.imshow, que é justamente usada para exibir a imagem).
Portanto, essa conversão não é necessária para o processamento na rede neural, mas é útil para a visualização e interpretação dos dados. Então, embora a imagem seja processada internamente como um vetor ao passar pelas camadas densas, ela é reformatada em uma matriz para facilitar a visualização.