Will we get an update on MPS for M2 Chip?

Wanted to know that will MPS work right off the shelf for the new M2 chip that Apple has just come out with? Or will we need to wait for an update on MPS to have support of it?

MPS should work right off the shelf. Apple has already implemented optimized kernels for the A15 (iPhone 13), which shares the same GPU architecture as the M2. This includes simd_shuffle_and_fill if it’s used at all in convolutions (probably not). Otherwise, it’s the exact same kernels as M1, using hardware-accelerated matrix multiplication via simdgroup_matrix. All of what I’m describing should be opaque to PyTorch, as the CPU-visible API of Metal Performance Shaders never changes across architectures or operating systems.