Continue local tensor mode enablement for DTensor tests

[dzmitry-huba]

Stack from ghstack (oldest at bottom):

• -> Continue local tensor mode enablement for DTensor tests #165451

cc @H-Huang @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @pragupta @msaroufim @dcci

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 the bot commands wiki

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

✅ You can merge normally! (1 Unrelated Failure)

As of commit d82fb80 with merge base 6adaa32 ():

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• trunk / macos-py3-arm64 / test (default, 1, 3, macos-m1-stable) (gh) (similar failure)
  RuntimeError: doctests 1/1 failed!

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

[dzmitry-huba]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

[albanD]

Thanks!
Could you also add a doc or update run_test.py to have instructions on how to run these tests?

[atalman]

hi @dzmitry-huba looks like this breaks torchao nightly builds since Oct 15 : https://github.com/pytorch/ao/actions/runs/18546536678/job/52865612926

I believe you need to do something like:

try:
    import numpy as np
    HAS_NUMPY = True
except ModuleNotFoundError:
    HAS_NUMPY = False
    np = None  # type: ignore[assignment]