Fix case when driver stack version is not available

Change-Id: Ie6e09982169cb8f435c2ba4683159deb2aa6fd20

python/tvm/relay/op/contrib/ethosn.py

@@ -47,8 +47,16 @@ def ethosn_available():
     return Available.SW_AND_HW if hw else Available.SW_ONLY
 
 
-def ethosn_api_version():
-    """Returns the version of the driver stack api that is being used."""
+def ethosn_api_version() -> str:
+    """
+    Returns the semantic version of the driver stack api that is
+    being used.
+
+    Returns
+    -------
+    str
+        Semantic version string (e.g. 3.0.1).
+    """
     return tvm.get_global_func("relay.ethos-n.api.version")()
 
 
@@ -85,7 +93,7 @@ def partition_for_ethosn(mod, params=None, **opts):
     expected_api_version = "3.0.1"
     if api_version != LooseVersion(expected_api_version):
         raise ValueError(
-            f"Driver stack version {api_version} unsupported. "
+            f"Driver stack version {api_version} is unsupported. "
             f"Please use version {expected_api_version}."
         )
 

tests/python/contrib/test_ethosn/test_split.py

@@ -17,15 +17,12 @@
 
 """Split tests for Arm(R) Ethos(TM)-N"""
 
-from distutils.version import LooseVersion
-
 import numpy as np
 import pytest
 
 import tvm
 from tvm import relay
 from tvm.testing import requires_ethosn
-from tvm.relay.op.contrib.ethosn import ethosn_api_version
 
 from . import infrastructure as tei
 
@@ -36,10 +33,7 @@ def _get_model(shape, dtype, splits, axis):
     return split.astuple()
 
 
-@pytest.mark.skipif(
-    ethosn_api_version() >= LooseVersion("3.0.1"),
-    reason="Split is not supported by this release of the driver stack",
-)
+@pytest.mark.skip("Split is not supported by the 3.0.1 version of the driver stack.")
 @requires_ethosn
 @pytest.mark.parametrize("dtype", ["uint8", "int8"])
 def test_split(dtype):
@@ -67,10 +61,7 @@ def test_split(dtype):
         tei.verify(outputs, dtype, 0)
 
 
-@pytest.mark.skipif(
-    ethosn_api_version() >= LooseVersion("3.0.1"),
-    reason="Split is not supported by this release of the driver stack",
-)
+@pytest.mark.skip("Split is not supported by the 3.0.1 version of the driver stack.")
 @requires_ethosn
 def test_split_failure():
     trials = [

tests/python/contrib/test_ethosn/test_topologies.py

@@ -17,15 +17,13 @@
 
 """Arm(R) Ethos(TM)-N tests for complex network topologies."""
 
-from distutils.version import LooseVersion
-
 import numpy as np
 import pytest
 
 import tvm
 from tvm import relay
 from tvm.testing import requires_ethosn
-from tvm.relay.op.contrib.ethosn import Available, ethosn_available, ethosn_api_version
+from tvm.relay.op.contrib.ethosn import Available, ethosn_available
 
 from . import infrastructure as tei
 
@@ -305,10 +303,7 @@ def get_model(shape, dtype, splits, axis):
             tei.verify(outputs, dtype, 0)
 
 
-@pytest.mark.skipif(
-    ethosn_api_version() >= LooseVersion("3.0.1"),
-    reason="Split is not supported by this release of the driver stack",
-)
+@pytest.mark.skip("Split is not supported by the 3.0.1 version of the driver stack.")
 @requires_ethosn
 @pytest.mark.parametrize("dtype", ["uint8", "int8"])
 def test_output_tuple_propagation(dtype):