More granularity in protecting tests requiring torch

Author: FrancescAlted

tests/ndarray/test_elementwise_funcs.py

@@ -6,8 +6,6 @@
 
 import blosc2
 
-torch = pytest.importorskip("torch", reason="torch not available")
-
 warnings.simplefilter("always")
 
 # Functions to test (add more as needed)
@@ -312,9 +310,10 @@ def test_unary_funcs(np_func, blosc_func, dtype, shape, chunkshape):
 @pytest.mark.parametrize(("np_func", "blosc_func"), UNARY_FUNC_PAIRS)
 @pytest.mark.parametrize("dtype", STR_DTYPES)
 @pytest.mark.parametrize("shape", [(10,), (20, 20)])
-@pytest.mark.parametrize("xp", [torch])
-def test_unfuncs_proxy(np_func, blosc_func, dtype, shape, xp):
-    _test_unary_func_proxy(np_func, blosc_func, dtype, shape, xp)
+def test_unary_funcs_torch_proxy(np_func, blosc_func, dtype, shape):
+    """Test unary functions with torch tensors as input (via proxy)."""
+    torch = pytest.importorskip("torch")
+    _test_unary_func_proxy(np_func, blosc_func, dtype, shape, torch)
 
 
 @pytest.mark.heavy
@@ -335,9 +334,10 @@ def test_binary_funcs(np_func, blosc_func, dtype, shape, chunkshape):
 @pytest.mark.parametrize(("np_func", "blosc_func"), BINARY_FUNC_PAIRS)
 @pytest.mark.parametrize("dtype", STR_DTYPES)
 @pytest.mark.parametrize(("shape", "chunkshape"), SHAPES_CHUNKS)
-@pytest.mark.parametrize("xp", [torch])
-def test_binfuncs_proxy(np_func, blosc_func, dtype, shape, chunkshape, xp):
-    _test_binary_func_proxy(np_func, blosc_func, dtype, shape, chunkshape, xp)
+def test_binary_funcs_torch_proxy(np_func, blosc_func, dtype, shape, chunkshape):
+    """Test binary functions with torch tensors as input (via proxy)."""
+    torch = pytest.importorskip("torch")
+    _test_binary_func_proxy(np_func, blosc_func, dtype, shape, chunkshape, torch)
 
 
 @pytest.mark.heavy

tests/ndarray/test_lazyexpr.py

@@ -15,7 +15,14 @@
 from blosc2.lazyexpr import ne_evaluate
 from blosc2.utils import get_chunks_idx, npvecdot
 
-torch = pytest.importorskip("torch", reason="torch not available")
+# Conditionally import torch for proxy tests
+try:
+    import torch
+
+    PROXY_TEST_XP = [torch, np]
+except ImportError:
+    torch = None
+    PROXY_TEST_XP = [np]
 
 NITEMS_SMALL = 100
 NITEMS = 1000
@@ -1848,7 +1855,7 @@ def test_lazyexpr_2args():
 
 @pytest.mark.parametrize(
     "xp",
-    [torch, np],
+    PROXY_TEST_XP,
 )
 @pytest.mark.parametrize(
     "dtype",

tests/ndarray/test_linalg.py

@@ -8,7 +8,14 @@
 from blosc2.lazyexpr import linalg_funcs
 from blosc2.utils import npvecdot
 
-torch = pytest.importorskip("torch", reason="torch not available")
+# Conditionally import torch for proxy tests
+try:
+    import torch
+
+    PROXY_TEST_XP = [torch, np]
+except ImportError:
+    torch = None
+    PROXY_TEST_XP = [np]
 
 
 @pytest.mark.parametrize(
@@ -827,7 +834,7 @@ def test_diagonal(shape, chunkshape, offset):
 
 @pytest.mark.parametrize(
     "xp",
-    [torch, np],
+    PROXY_TEST_XP,
 )
 @pytest.mark.parametrize(
     "dtype",

tests/ndarray/test_setitem.py

@@ -11,8 +11,6 @@
 
 import blosc2
 
-torch = pytest.importorskip("torch", reason="torch not available")
-
 argnames = "shape, chunks, blocks, slices, dtype"
 argvalues = [
     ([456], [258], [73], slice(0, 1), np.int32),
@@ -46,14 +44,6 @@ def test_setitem(shape, chunks, blocks, slices, dtype):
     nparray[slices] = val
     np.testing.assert_almost_equal(a[...], nparray)
 
-    # Object called via SimpleProxy
-    slice_shape = a[slices].shape
-    dtype_ = {np.float32: torch.float32, np.int32: torch.int32, np.float64: torch.float64}[dtype]
-    val = torch.ones(slice_shape, dtype=dtype_)
-    a[slices] = val
-    nparray[slices] = val
-    np.testing.assert_almost_equal(a[...], nparray)
-
     # blosc2.NDArray
     if np.prod(slice_shape) == 1 or len(slice_shape) != len(blocks):
         chunks = None
@@ -65,6 +55,22 @@ def test_setitem(shape, chunks, blocks, slices, dtype):
     np.testing.assert_almost_equal(a[...], nparray)
 
 
+@pytest.mark.parametrize(argnames, argvalues)
+def test_setitem_torch_proxy(shape, chunks, blocks, slices, dtype):
+    torch = pytest.importorskip("torch")
+    size = int(np.prod(shape))
+    nparray = np.arange(size, dtype=dtype).reshape(shape)
+    a = blosc2.frombuffer(bytes(nparray), nparray.shape, dtype=dtype, chunks=chunks, blocks=blocks)
+
+    # Object called via SimpleProxy (torch tensor)
+    slice_shape = a[slices].shape
+    dtype_ = {np.float32: torch.float32, np.int32: torch.int32, np.float64: torch.float64}[dtype]
+    val = torch.ones(slice_shape, dtype=dtype_)
+    a[slices] = val
+    nparray[slices] = val
+    np.testing.assert_almost_equal(a[...], nparray)
+
+
 @pytest.mark.parametrize(
     ("shape", "slices"),
     [