MPS Support

[Mahnoo06]

Hello, I am using functorch for my project. It's a good substitute for Jax and makes life easy. However it does not have GPU support for M1 chips which is a time-constraint for me because I have an M1-based system. Now that PyTorch support M1-GPU I think functorch should also make this feature available.

[zou3519]

Hey @Mahnoo06, MPS support is a "prototype" feature in PyTorch so we haven't actively tried to support it yet. That being said I know M1 machines are popular, so we'll try to enable this.