Skip test if import missing

parsiad · Aug 23, 2024 · 63f54c0 · 63f54c0
1 parent 694ccdb
commit 63f54c0
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/tests/test_func.py b/tests/test_func.py
@@ -1,10 +1,13 @@
 import numpy as np
-import scipy.special
+import pytest
 
 import micrograd_pp as mpp
 
 
+@pytest.mark.skipif(not pytest.importorskip("scipy.special"), reason="Unable to import scipy.special")
 def test_softmax() -> None:
+    import scipy.special
+
     a = np.random.randn(5, 4, 3)
     actual = mpp.softmax(mpp.Constant(a), dim=1).value
     desired = scipy.special.softmax(a, axis=1)

diff --git a/tests/test_nn.py b/tests/test_nn.py
@@ -2,7 +2,6 @@
 
 import numpy as np
 import pytest
-import torch
 
 import micrograd_pp as mpp
 
@@ -88,7 +87,10 @@ def test_layer_norm() -> None:
 
 
 @pytest.mark.parametrize("is_causal", (False, True))
+@pytest.mark.skipif(not pytest.importorskip("torch"), reason="Unable to import torch")
 def test_multihead_attention(is_causal: bool) -> None:  # Test against PyTorch implementation
+    import torch
+
     torch_attn_mask = None
     mpp_attn_mask = None
     if is_causal: