[FRONTEND] Add check for the axis of reduction op (triton-lang#1268)

siliconflow · Mar 7, 2023 · e100f8d · e100f8d
1 parent a89586f
commit e100f8d
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/python/triton/language/semantic.py b/python/triton/language/semantic.py
@@ -1133,6 +1133,10 @@ def reduce_impl(input: tl.tensor, axis: int, builder: ir.builder, name: str,
 
     # get result type
     shape = input.type.shape
+
+    rank = len(shape)
+    assert 0 <= axis < rank, f"axis (v={axis}) is out of range, should be within [0, {rank})"
+
     ret_shape = []
     for i, s in enumerate(shape):
         if i != axis: