Moving CommTensor from tests to private _spmd folder

[mrshenli]

Stack from ghstack (oldest at bottom):

• -> Moving CommTensor from tests to private _spmd folder #84719

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/84719

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures, 3 Pending

As of commit c80ce69:
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[mrshenli]

Compared to #84655, the fix is not exposing anything in _spmd/__init__.py file, otherwise, the imported non-public APIs in comm_tensor.py will break test_public_bindings.py.

[wanchaol]

Compared to #84655, the fix is not exposing anything in _spmd/__init__.py file, otherwise, the imported non-public APIs in comm_tensor.py will break test_public_bindings.py.

hmmm I didn't get this fully, why test_public_bindings would be bothered for anything inside a private module?

[mrshenli]

hmmm I didn't get this fully, why test_public_bindings would be bothered for anything inside a private module?

Error message looks like this. Before this PR, experimental's __init__.py didn't expose anything to that test. However, if we expose CommTensor which imports those private features, it makes those private features visible to the test which complains about missing __all__ in the experimental package.

Full list:
# torch.fx.experimental.proxy_tensor.ProxySymDispatchMode:
  - Is NOT public: it is not inside the module's (`torch.fx.experimental.proxy_tensor`) `__all__`
  - Does look public: it does look public because it follows the rules from the doc above (does not start with `_` and has a proper `__module__`).
  - You can do either of these two things to fix this problem:
    - To make it public: add it from the modules's (`torch.fx.experimental.proxy_tensor`) `__all__`
    - To make it NOT look public: make its name start with `_`

[mrshenli]

@pytorchbot merge -g

[pytorchmergebot]

@pytorchbot successfully started a merge job. Check the current status here.
The merge job was triggered with the green (-g) flag. This means that your change will be merged once all checks on your PR have passed (ETA: 0-4 Hours). If this is not the intended behavior, feel free to use some of the other merge options in the wiki.
Please reach out to the PyTorch DevX Team with feedback or questions!

[github-actions]

Hey @mrshenli.
You've committed this PR, but it does not have both a 'release notes: ...' and 'topics: ...' label. Please add one of each to the PR. The 'release notes: ...' label should represent the part of PyTorch that this PR changes (fx, autograd, distributed, etc) and the 'topics: ...' label should represent the kind of PR it is (not user facing, new feature, bug fix, perf improvement, etc). The list of valid labels can be found here for the 'release notes: ...' and here for the 'topics: ...'.
For changes that are 'topic: not user facing' there is no need for a release notes label.