Adds pytorch to get_namespace

Author: thomasjpfan

array_api_compat/common/_helpers.py

@@ -75,6 +75,13 @@ def get_namespace(*xs, _use_compat=True):
             else:
                 import cupy as cp
                 namespaces.add(cp)
+        elif _is_torch_array(x):
+            if _use_compat:
+                from .. import torch as torch_namespace
+                namespaces.add(torch_namespace)
+            else:
+                import torch
+                namespaces.add(torch)
         else:
             # TODO: Support Python scalars?
             raise ValueError("The input is not a supported array type")

tests/test_get_namespace.py

@@ -0,0 +1,14 @@
+import array_api_compat
+import pytest
+
+
+@pytest.mark.parametrize("library", ["cupy", "numpy", "torch"])
+def test_get_namespace(library):
+    lib = pytest.importorskip(library)
+
+    array = lib.asarray([1.0, 2.0, 3.0])
+    namespace = array_api_compat.get_namespace(array)
+
+    expected_namespace = getattr(array_api_compat, library)
+    assert namespace is expected_namespace
+