[MPS] Extend autocast support to bf16 #139386
Description
🚀 The feature, motivation and pitch
In #99272, autocast support was added for fp16 (torch.float16) in MPS. We would like to extend this functionality to include bf16 (torch.bfloat16). Most of the code required for this feature can be referenced in #99272, so please review prior commits for guidance.
Key Requirements
- macOS Support: bf16 support is only available on macOS 14 and later; the feature should include checks to ensure compatibility.
- Error Handling: Implement clear error messages when bf16 is used on unsupported macOS versions.
- Testing: Add comprehensive unit tests for bf16 autocast, with emphasis on edge cases.
- Documentation: Update relevant documentation to include bf16 support details, including version requirements, usage examples, and potential performance considerations (if applicable).
Alternatives
No response
Additional context
No response