MAINT: run self-tests even if a library is missing

Author: ev-br

tests/test_array_namespace.py

@@ -2,10 +2,8 @@
 import sys
 import warnings
 
-import jax
 import numpy as np
 import pytest
-import torch
 
 import array_api_compat
 from array_api_compat import array_namespace
@@ -76,6 +74,7 @@ def test_array_namespace(library, api_version, use_compat):
         subprocess.run([sys.executable, "-c", code], check=True)
 
 def test_jax_zero_gradient():
+    jax = import_("jax")
     jx = jax.numpy.arange(4)
     jax_zero = jax.vmap(jax.grad(jax.numpy.float32, allow_int=True))(jx)
     assert array_namespace(jax_zero) is array_namespace(jx)
@@ -89,11 +88,13 @@ def test_array_namespace_errors():
     pytest.raises(TypeError, lambda: array_namespace(x, (x, x)))
 
 def test_array_namespace_errors_torch():
+    torch = import_("torch")
     y = torch.asarray([1, 2])
     x = np.asarray([1, 2])
     pytest.raises(TypeError, lambda: array_namespace(x, y))
 
 def test_api_version_torch():
+    torch = import_("torch")
     x = torch.asarray([1, 2])
     torch_ = import_("torch", wrapper=True)
     assert array_namespace(x, api_version="2023.12") == torch_
@@ -118,6 +119,7 @@ def test_get_namespace():
     assert array_api_compat.get_namespace is array_namespace
 
 def test_python_scalars():
+    torch = import_("torch")
     a = torch.asarray([1, 2])
     xp = import_("torch", wrapper=True)
 

tests/test_dask.py

@@ -1,10 +1,14 @@
 from contextlib import contextmanager
 
 import array_api_strict
-import dask
 import numpy as np
 import pytest
-import dask.array as da
+
+try:
+    import dask
+    import dask.array as da
+except ImportError:
+    pytestmark = pytest.skip(allow_module_level=True, reason="dask not found")
 
 from array_api_compat import array_namespace
 

tests/test_jax.py

@@ -1,10 +1,14 @@
-import jax
-import jax.numpy as jnp
 from numpy.testing import assert_equal
 import pytest
 
 from array_api_compat import device, to_device
 
+try:
+    import jax
+    import jax.numpy as jnp
+except ImportError:
+    pytestmark = pytest.skip(allow_module_level=True, reason="jax not found")
+
 HAS_JAX_0_4_31 = jax.__version__ >= "0.4.31"
 
 

tests/test_vendoring.py

@@ -16,11 +16,13 @@ def test_vendoring_cupy():
 
 
 def test_vendoring_torch():
+    pytest.importorskip("torch")
     from vendor_test import uses_torch
 
     uses_torch._test_torch()
 
 
 def test_vendoring_dask():
+    pytest.importorskip("dask")
     from vendor_test import uses_dask
     uses_dask._test_dask()