docarray
diff --git a/‎tests/units/array/stack/test_array_stacked_jax.py‎
Lines changed: 5 additions & 1 deletion b/‎tests/units/array/stack/test_array_stacked_jax.py‎
Lines changed: 5 additions & 1 deletion
diff --git a/‎tests/units/computation_backends/jax_backend/test_basics.py‎
Lines changed: 2 additions & 11 deletions b/‎tests/units/computation_backends/jax_backend/test_basics.py‎
Lines changed: 2 additions & 11 deletions
@@ -1,6 +1,5 @@
 from typing import Optional, Union
 
-import jax.numpy as jnp
 import pytest
 
 from docarray import BaseDoc, DocList
@@ -14,6 +13,11 @@
     NdArray,
     VideoTensor,
 )
+from docarray.utils._internal.misc import is_jax_available
+
+jax_available = is_jax_available()
+if jax_available:
+    import jax.numpy as jnp
 
 
 @pytest.fixture()
 
@@ -9,7 +9,8 @@
 
     from docarray.computation.jax_backend import JaxCompBackend
     from docarray.typing import JaxArray
-jax.config.update("jax_enable_x64", True)
+
+    jax.config.update("jax_enable_x64", True)
 
 
 @pytest.mark.parametrize(
@@ -26,7 +27,6 @@ def test_n_dim(shape, result):
     assert JaxCompBackend.n_dim(array) == result
 
 
-@pytest.mark.tensorflow
 @pytest.mark.parametrize(
     'shape,result',
     [
@@ -42,14 +42,12 @@ def test_shape(shape, result):
     assert type(shape) == tuple
 
 
-@pytest.mark.tensorflow
 def test_to_device():
     array = JaxArray(jnp.zeros((3)))
     array = JaxCompBackend.to_device(array, 'cpu')
     assert array.tensor.device().platform.endswith('cpu')
 
 
-@pytest.mark.tensorflow
 @pytest.mark.parametrize(
     'dtype,result_type',
     [
@@ -64,34 +62,29 @@ def test_dtype(dtype, result_type):
     assert JaxCompBackend.dtype(array) == result_type
 
 
-@pytest.mark.tensorflow
 def test_empty():
     array = JaxCompBackend.empty((10, 3))
     assert array.tensor.shape == (10, 3)
 
 
-@pytest.mark.tensorflow
 def test_empty_dtype():
     tf_tensor = JaxCompBackend.empty((10, 3), dtype=jnp.int32)
     assert tf_tensor.tensor.shape == (10, 3)
     assert tf_tensor.tensor.dtype == jnp.int32
 
 
-@pytest.mark.tensorflow
 def test_empty_device():
     tensor = JaxCompBackend.empty((10, 3), device='cpu')
     assert tensor.tensor.shape == (10, 3)
     assert tensor.tensor.device().platform.endswith('cpu')
 
 
-@pytest.mark.tensorflow
 def test_squeeze():
     tensor = JaxArray(jnp.zeros(shape=(1, 1, 3, 1)))
     squeezed = JaxCompBackend.squeeze(tensor)
     assert squeezed.tensor.shape == (3,)
 
 
-@pytest.mark.tensorflow
 @pytest.mark.parametrize(
     'data_input,t_range,x_range,data_result',
     [
@@ -123,14 +116,12 @@ def test_minmax_normalize(data_input, t_range, x_range, data_result):
     assert jnp.allclose(output.tensor, jnp.array(data_result))
 
 
-@pytest.mark.tensorflow
 def test_reshape():
     tensor = JaxArray(jnp.zeros((3, 224, 224)))
     reshaped = JaxCompBackend.reshape(tensor, (224, 224, 3))
     assert reshaped.tensor.shape == (224, 224, 3)
 
 
-@pytest.mark.tensorflow
 def test_stack():
     t0 = JaxArray(jnp.zeros((3, 224, 224)))
     t1 = JaxArray(jnp.ones((3, 224, 224)))