Optimized the test cases

Author: git-flyer

python/test/tle/unit/test_tle_distributed.py

@@ -228,11 +228,12 @@ def __init__(self):
 
 class TestShardId:
 
-    def test_node_axis_requires_device_dptr(self):
+    @pytest.mark.parametrize("axis", ("device", "node"))
+    def test_rank_axis_requires_device_dptr(self, axis):
         mesh = tle.device_mesh({"node": 2, "device": 4})
         semantic = _FakeSemantic()
-        with pytest.raises(ValueError, match="device_dptr is required for axis 'node'"):
-            tle.shard_id(mesh, "node", _semantic=semantic)
+        with pytest.raises(ValueError, match=rf"device_dptr is required for axis '{axis}'"):
+            tle.shard_id(mesh, axis, _semantic=semantic)
 
 
 class TestDistributedBarrierScope:

python/triton/experimental/tle/language/distributed.py

@@ -656,7 +656,7 @@ def shard_id(
     mesh = tl._unwrap_if_constexpr(mesh)
     axis = tl._unwrap_if_constexpr(axis)
 
-    if axis == "node" and device_dptr is None:
+    if axis in ("device", "node") and device_dptr is None:
         raise ValueError(f"device_dptr is required for axis {axis!r}")
 
     if axis == "device":