mirror of
https://github.com/zebrajr/pytorch.git
synced 2026-01-15 12:15:51 +00:00
452abac61c
**Summary:** Currently, whenever we subtract two partial dtensors, we redistribute since linearity is -1 for aten.sub.tensor. However, this is an unnecessary redistribution that can be avoided in similar ways to its add counterpart. I moved the op to linear_ops and ensured subtracting a scalar from a partial dtensor continues to redistribute. **Test Cases:** 1. pytest test/distributed/tensor/test_pointwise_ops.py -k test_add_sub_scalar_norm_partial 2. pytest test/distributed/tensor/test_pointwise_ops.py -k test_add_sub_scalar_partial Pull Request resolved: https://github.com/pytorch/pytorch/pull/170040 Approved by: https://github.com/wconstab ghstack dependencies: #170030, #170035