[FRONTEND] don't call tl.view in arg{min,max} (triton-lang#1518)

A small oversight in triton-lang#1305, since `view` can rearrange elements it should be avoided here. Instead I use indexing with `None` to create new dimensions. Co-authored-by: Philippe Tillet <[email protected]>
siliconflow · Apr 13, 2023 · db9a1cd · db9a1cd
1 parent c5359f0
commit db9a1cd
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/python/triton/language/core.py b/python/triton/language/core.py
@@ -1255,9 +1255,10 @@ def _argreduce(input, axis, combine_fn, _builder=None, _generator=None):
     index = arange(0, n, _builder=_builder)
 
     if len(input.shape) > 1:
-        new_shape = [constexpr(1)] * len(input.shape)
-        new_shape[axis] = constexpr(n)
-        index = view(index, new_shape, _builder=_builder)
+        # Broadcast index across the non-reduced axes
+        expand_dims_index = [None] * len(input.shape)
+        expand_dims_index[axis] = slice(None)
+        index = index.__getitem__(expand_dims_index, _builder=_builder)
         index = broadcast_to(index, input.shape, _builder=_builder)
 
     rvalue, rindices = reduction((input, index), axis, combine_fn,