[BE][Easy] export explicitly imported public submodules

[XuehaiPan]

Stack from ghstack (oldest at bottom):

• [BE] enable UFMT in torch.utils.data #127705
• [BE] sort imports in torch.utils.data #127704
• [BE] explicitly export subpackage torch.utils #128342
• [BE] enable UFMT for top-level files torch/*.py #127707
• [BE] enable UFMT for torch/storage.py #127706
• [BE] enable UFMT for torch/__init__.py #127710
• [BE][Easy] sort __all__ in torch/__init__.py #127709
• [BE] sort imports in torch/__init__.py #127708
• -> [BE][Easy] export explicitly imported public submodules #127703

Add top-level submodules torch.{storage,serialization,functional,amp,overrides,types}

cc @albanD @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @chauhang

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 or our office hours

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

✅ You can merge normally! (3 Unrelated Failures)

As of commit d7fd93d with merge base 75b0720 ():

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

• pull / linux-focal-py3_8-clang9-xla / test (xla, 1, 1, linux.12xlarge) (gh) (similar failure)
  test_all_reduce_no_op_with_one_replica
• trunk / linux-focal-cuda12.4-py3.10-gcc9-sm86 / test (default, 3, 5, linux.g5.4xlarge.nvidia.gpu) (gh) (similar failure)
  test_foreach.py::TestForeachCUDA::test_binary_op_list_error_cases__foreach_add_cuda_bfloat16

UNSTABLE - The following job failed but was likely due to flakiness present on trunk and has been marked as unstable:

• inductor / linux-jammy-cpu-py3.8-gcc11-inductor / test (inductor_torchbench_cpu_smoketest_perf, 1, 1, linux.24xl.spr-metal, unstable) (gh)
  Process completed with exit code 1.

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

[Skylion007]

Some of these changes seem related to: #127688 and probably should wait on that.

[Skylion007]

Is ufmt forcing the removal of all these relative imports?

[XuehaiPan]

No. This is updated by hand.

Absolute imports, or relative imports from siblings, are recommended by PEP8.

Absolute imports are recommended, as they are usually more readable and tend to be better behaved (or at least give better error messages) if the import system is incorrectly configured (such as when a directory inside a package ends up on sys.path):

import mypkg.sibling
from mypkg import sibling
from mypkg.sibling import example

However, explicit relative imports are an acceptable alternative to absolute imports, especially when dealing with complex package layouts where using absolute imports would be unnecessarily verbose:

from . import sibling
from .sibling import example

Standard library code should avoid complex package layouts and always use absolute imports.

---

We can enforce this by adding the ruff rule TID (flake8-tidy-imports) with:

[tool.ruff.lint.flake8-tidy-imports]
ban-relative-imports = "all"  # defaults to "parents"; see https://docs.astral.sh/ruff/settings/#lintflake8-tidy-imports

[Skylion007]

Can't this affect import ordering which can be bad for our Python submodules in torch that do have side effects (or cause cyclic import errors where non existed before).

[XuehaiPan]

Changing relative import to identical absolute import does not change the side effects of the import statement.

[albanD]

Doesn't it change the import order though?
Where the relative import would actually trigger the import of the file while the absolute import would expect the partially initialized module to already contain the corresponding entry?

[XuehaiPan]

This change is made based on:

1. The relative import will convert to absolute import in bytecode eval loop:

$ echo 'from torch.optim import adam' | python3 -m dis -
  0           0 RESUME                   0

  1           2 LOAD_CONST               0 (0)
              4 LOAD_CONST               1 (('adam',))
              6 IMPORT_NAME              0 (torch.optim)
              8 IMPORT_FROM              1 (adam)
             10 STORE_NAME               1 (adam)
             12 POP_TOP
             14 RETURN_CONST             2 (None)
$ echo 'from . import adam' | python3 -m dis -
  0           0 RESUME                   0

  1           2 LOAD_CONST               0 (1)
              4 LOAD_CONST               1 (('adam',))
              6 IMPORT_NAME              0
              8 IMPORT_FROM              1 (adam)
             10 STORE_NAME               1 (adam)
             12 POP_TOP
             14 RETURN_CONST             2 (None)

https://github.com/python/cpython/blob/a9f2daf1ab182d95b44ee94dc9fb8faec60e34b1/Python/ceval.c#L2513-L2526

from . import adam is convert to __import__(f'{__name__}.adam') where __name__ = 'torch.optim'.

They are semantically identical.

2. See [BE][Easy] export explicitly imported public submodules #127703 (comment): Absolute imports, or relative imports from siblings, are recommended by PEP8.

[fbgheith]

How about absolute paths when combined with circular dependencies? Example:

torch/optim/init.py import torch.optim._multi_tensor and
torch/optim/_multi_tenoser/init.py import torch.optim

This pattern, combined with absolute import paths is causing unpredictable and hard to debug failures.

[XuehaiPan]

This pattern, combined with absolute import paths is causing unpredictable and hard to debug failures.

You can add a raise ImportError statement at the top of the file, when looking at the traceback of python3 -c 'import torch'. This can be used to find the first dependent of the submodule.

[malfet]

@fbgheith reports that this broke some lazy import logic (still trying
to figure out the details), so I'll be attempting to revert this and few others that depend on it to restore original behavior)

[XuehaiPan]

@fbgheith reports that this broke some lazy import logic (still trying to figure out the details), so I'll be attempting to revert this and few others that depend on it to restore original behavior)

I will help with fixing that in PR #129095.

• re-export torch.optim._multi_tensor in torch/__init__.py #129095

[fbgheith]

This has been a big source of headache where depending on the exact ordering if imports in the pytorch client code one of three things would happen:

(1) great majority: things work fine
(2a) the "import torch.optim._multi_tensor" in torch/optim/init.py causes infinite recursion in a small set of uses cases
(2b) removing the import above solves (2a) but makes "import torch.optim._multi_tensor" fail for another small set
(2c) replacing the backwards absolute import with relative ones make both sets happy

This is just empirical data for this pair of dependencies.

[fbgheith]

Please keep in mind that #129095 is just a bandaid to remedy the particular scenario described above without reverting #127703. We also need to come up with an explanation for why we ended up in this state and the proper guidelines for what can go in init.py files.

[XuehaiPan]

We also need to come up with an explanation for why we ended up in this state and the proper guidelines for what can go in __init__.py files.

We will be able to do this after we enable ufmt globally and optionally enable ruff's TID rules.