fixes

Author: crusaderky

array_api_compat/torch/__init__.py

@@ -3,21 +3,25 @@
 from torch import * # noqa: F403
 
 # Several names are not included in the above import *
-_torch_all = set()
+_torch_dir = set()
 import torch
 for n in dir(torch):
     if (n.startswith('_')
         or n.endswith('_')
-        or 'cuda' in n
-        or 'cpu' in n
         or 'backward' in n):
         continue
     exec(f"{n} = torch.{n}")
-    _torch_all.add(n)
+    _torch_dir.add(n)
 del n
 
+# torch.__all__ is wildly incorrect
+_n: dict[str, object] = {}
+exec('from torch import *', _n)
+_torch_all = set(_n)
+del _n
+
 # These imports may overwrite names from the import * above.
-import _aliases
+from . import _aliases
 from ._aliases import * # noqa: F403
 from ._info import __array_namespace_info__  # noqa: F401
 
@@ -31,7 +35,8 @@
     set(_torch_all)
     | set(_aliases.__all__)
     | {"__array_api_version__", "__array_namespace_info__", "linalg", "fft"}
+    | {"from_dlpack"}
 )
 
 def __dir__() -> list[str]:
-    return __all__
+    return sorted(set(__all__) | set(_torch_dir))

array_api_compat/torch/fft.py

@@ -73,11 +73,15 @@ def ifftshift(
     return torch.fft.ifftshift(x, dim=axes, **kwargs)
 
 
-__all__ = torch.fft.__all__ + [
+_all = {
     "fftn",
     "ifftn",
     "rfftn",
     "irfftn",
     "fftshift",
     "ifftshift",
-]
+}
+__all__ = sorted(set(torch.fft.__all__) |_all)
+
+def __dir__() -> list[str]:
+    return sorted(set(dir(torch.fft)) | _all)

tests/test_all.py

@@ -12,6 +12,9 @@
         # Inspection
         "__array_api_version__",
         "__array_namespace_info__",
+        # Submodules
+        "fft",
+        "linalg",
         # Constants
         "e",
         "inf",
@@ -240,6 +243,8 @@ def all_names(mod):
     This is typically `__all__` but, if not defined, Python
     implements automated fallbacks.
     """
+    # Note: this method also makes the test trip if a name is
+    # in __all__ but doesn't actually appear in the module.
     objs = {}
     exec(f"from {mod.__name__} import *", objs)
     return list(objs)