fix(ascend): correct flattened argmin semantics

argmin(inp) with dim=None returned wrong indices on the Ascend backend.
Three issues on the same code path:

1. Tie breaking. tl.min(..., return_indices=True) did not request
   return_indices_tie_break_left, so equal values could resolve to an
   arbitrary index instead of the first one. PyTorch returns the first
   occurrence.
2. Block size. block_size was derived only from sqrt(numel) with no upper
   bound. For large inputs it exceeded what the device could handle on this
   path and produced wrong indices. It is now capped at 1024.
3. keepdim. The output was allocated with torch.empty([]) unconditionally,
   so keepdim=True was dropped. PyTorch returns shape [1] * inp.dim().

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: zheng1

src/flag_gems/runtime/backend/_ascend/ops/argmin.py

@@ -43,7 +43,12 @@ def argmin_kernel_1(
 
     max_value = get_dtype_max(inp.type.element_ty)
     inp_val = tl.load(inp_ptrs, mask=mask, other=max_value)
-    min_val, min_index = tl.min(inp_val, axis=0, return_indices=True)
+    min_val, min_index = tl.min(
+        inp_val,
+        axis=0,
+        return_indices=True,
+        return_indices_tie_break_left=True,
+    )
     min_index = min_index + pid * BLOCK_SIZE
     mid_value_ptr = mid_value + pid
     min_index_ptr = mid_index + pid
@@ -128,13 +133,14 @@ def argmin(inp, dim=None, keepdim=False, *, dtype=None):
         M = inp.numel()
         if dtype is None:
             dtype = inp.dtype
-        block_size = triton.next_power_of_2(math.ceil(math.sqrt(M)))
+        block_size = min(1024, triton.next_power_of_2(math.ceil(math.sqrt(M))))
         mid_size = triton.cdiv(M, block_size)
         block_mid = triton.next_power_of_2(mid_size)
 
         mid_value = torch.empty((mid_size,), dtype=dtype, device=inp.device)
         mid_index = torch.empty((mid_size,), dtype=torch.int64, device=inp.device)
-        out = torch.empty([], dtype=torch.int64, device=inp.device)
+        out_shape = [1] * inp.dim() if keepdim else []
+        out = torch.empty(out_shape, dtype=torch.int64, device=inp.device)
 
         with torch_device_fn.device(inp.device):
             argmin_kernel_1[(mid_size, 1, 1)](

tests/test_argmin.py

@@ -39,10 +39,42 @@ def test_argmin(shape, dim, keepdim, dtype):
     else:
         inp = torch.randn(shape, dtype=dtype, device=flag_gems.device)
 
-    ref_inp = utils.to_reference(inp)
+    use_cpu_ref = flag_gems.vendor_name == "ascend" and dim is None and keepdim
+    ref_inp = inp.cpu() if use_cpu_ref else utils.to_reference(inp)
     ref_out = torch.argmin(ref_inp, dim=dim, keepdim=keepdim)
+    if use_cpu_ref and not cfg.TO_CPU:
+        ref_out = ref_out.to(inp.device)
 
     with flag_gems.use_gems():
         res_out = torch.argmin(inp, dim=dim, keepdim=keepdim)
 
     utils.gems_assert_equal(res_out, ref_out)
+
+
+@pytest.mark.argmin
+@pytest.mark.skipif(
+    flag_gems.vendor_name != "ascend", reason="regression test for Ascend"
+)
+def test_argmin_large_flattened_index():
+    inp = torch.ones((200, 2560, 3), dtype=torch.float32, device=flag_gems.device)
+    inp.flatten()[1024 * 1200] = -1
+    ref_out = torch.argmin(inp.cpu())
+
+    with flag_gems.use_gems():
+        res_out = torch.argmin(inp)
+
+    torch.testing.assert_close(res_out.cpu(), ref_out, atol=0, rtol=0)
+
+
+@pytest.mark.argmin
+@pytest.mark.skipif(
+    flag_gems.vendor_name != "ascend", reason="regression test for Ascend"
+)
+def test_argmin_flattened_nan_index():
+    inp = torch.tensor([3.0, float("nan"), -5.0, float("nan")], device=flag_gems.device)
+    ref_out = torch.argmin(inp.cpu())
+
+    with flag_gems.use_gems():
+        res_out = torch.argmin(inp)
+
+    torch.testing.assert_close(res_out.cpu(), ref_out, atol=0, rtol=0)