test_argmin.py

# Copyright 2026 FlagOS Contributors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import pytest
import torch

import flag_gems

from . import accuracy_utils as utils
from . import conftest as cfg

if cfg.QUICK_MODE:
    FLOAT_DTYPES = [torch.float32]
    DIM_LIST = [1]
else:
    FLOAT_DTYPES = utils.FLOAT_DTYPES
    DIM_LIST = [0, 1]


@pytest.mark.argmin
@pytest.mark.parametrize("shape", utils.REDUCTION_SMALL_SHAPES)
@pytest.mark.parametrize("dim", DIM_LIST + [None])
@pytest.mark.parametrize("keepdim", [True, False])
@pytest.mark.parametrize("dtype", FLOAT_DTYPES + utils.INT_DTYPES)
def test_argmin(shape, dim, keepdim, dtype):
    if dtype in utils.INT_DTYPES:
        inp = torch.randint(-1024, 1024, size=shape, device=flag_gems.device).to(dtype)
    else:
        inp = torch.randn(shape, dtype=dtype, device=flag_gems.device)

    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)