xla
3632f16e - [Functionalization] Enable FSDP

Commit
2 years ago
[Functionalization] Enable FSDP Summary: This pull request enables FSDP by replacing .set_ with our own _replace_xla_tensor API. The reason for that is Functionalization pass will reapply the new value to all the tensor's aliases since it's an in-place ops. However, that reapplication assumes the source and the destination would share the same amount of elements (view_copy). And .set_ doesn't follow this rule. P.S. It also removes two .data tests that are no longer applicable. Test Plan: CI.
Author
Committer
Parents
Loading