Add missing names to torch.linalg

Author: asmeurer

array_api_compat/cupy/__init__.py

@@ -6,13 +6,7 @@
 # These imports may overwrite names from the import * above.
 from ._aliases import *
 
-# Don't know why, but we have to do an absolute import to import linalg. If we
-# instead do
-#
-# from . import linalg
-#
-# It doesn't overwrite cupy.linalg from above. The import is generated
-# dynamically so that the library can be vendored.
+# See the comment in the numpy __init__.py
 __import__(__package__ + '.linalg')
 
 from .linalg import matrix_transpose, vecdot

array_api_compat/torch/__init__.py

@@ -14,6 +14,9 @@
 # These imports may overwrite names from the import * above.
 from ._aliases import *
 
+# See the comment in the numpy __init__.py
+__import__(__package__ + '.linalg')
+
 from ..common._helpers import *
 
 __array_api_version__ = '2021.12'

array_api_compat/torch/linalg.py

@@ -1 +1,15 @@
-raise ImportError("The array api compat torch.linalg module extension is not yet implemented")
+from torch.linalg import *
+
+# torch.linalg doesn't define __all__
+# from torch.linalg import __all__ as linalg_all
+from torch import linalg as _linalg
+linalg_all = [i for i in dir(_linalg) if not i.startswith('_')]
+
+# These are implemented in torch but aren't in the linalg namespace
+from torch import outer, trace
+from ._aliases import matrix_transpose, tensordot
+
+__all__ = linalg_all + ['outer', 'trace', 'matrix_transpose', 'tensordot']
+
+del linalg_all
+del _linalg