diff --git a/coreai_torch/_aten_to_core.py b/coreai_torch/_aten_to_core.py index cbad1d2..dac2c37 100644 --- a/coreai_torch/_aten_to_core.py +++ b/coreai_torch/_aten_to_core.py @@ -2311,11 +2311,22 @@ def replace_min_dim( dim = dim + x.type.rank if dim < 0 else dim min_values = coreai.reduce_min(x, [dim]) + + element_type = x.type.element_type + is_integer = isinstance(element_type, IntegerType) + is_bool = element_type == IntegerType.get_signless(1) + if is_integer and not is_bool: + # ~x == x ^ ALL_ONES; -1 has all bits set in two's complement. Bitwise + # complement is a strictly order-reversing bijection over the whole + # integer range (no overflow), so argmax(~x) == argmin(x). + reversed_x = coreai.broadcasting_bitwise_xor( + x, coreai.constant(-1, dtype=element_type) + ) + else: + reversed_x = coreai.broadcasting_mul(x, coreai.constant(-1, dtype=element_type)) + argmin_indices = coreai.cast( - coreai.argmax( - coreai.broadcasting_mul(x, coreai.constant(-1, dtype=x.type.element_type)), - dim, - ), + coreai.argmax(reversed_x, dim), np.int32, ) diff --git a/tests/ops/test_ops.py b/tests/ops/test_ops.py index 4c24e96..da180d0 100644 --- a/tests/ops/test_ops.py +++ b/tests/ops/test_ops.py @@ -3003,6 +3003,86 @@ def forward(self, x: Tensor) -> tuple[Tensor, Tensor]: await validate_numerical_output(model=model, x=x, dynamic_shapes=dynamic_shapes) +class TestMinDimArgminDtypeExtremes: + """min.dim argmin indices must be correct even at dtype-extremal minima. + + Core AI has argmax but not argmin, so ``replace_min_dim`` derives the + argmin index by reversing the order of ``x`` and taking ``argmax``. The + reversal must be exact at the extremes of the input dtype's range: + + * uint8 minimum ``0`` — a plain ``-x`` (i.e. ``x * -1``) negation maps + ``0`` to ``0`` rather than the top of the unsigned range, so the true + minimum stops being the argmax. + * int8 minimum ``-128`` — ``-128 * -1`` overflows the signed range and + (under wrapping arithmetic) lands back on ``-128``. + + Both were previously wrong on interpreter/cpu/gpu. The fix reverses integer + inputs with the bitwise complement ``~x`` (a strictly decreasing bijection + over the whole integer range that never overflows), leaving the exact float + negation path untouched. Ground truth is ``torch.min(x, dim=...)``. + """ + + class MinDimModel(nn.Module): + def __init__(self, dim: int, keepdim: bool) -> None: + super().__init__() + self.dim = dim + self.keepdim = keepdim + + def forward(self, x: Tensor) -> tuple[Tensor, Tensor]: + return torch.min(x, dim=self.dim, keepdim=self.keepdim) + + @pytest.mark.parametrize("keepdim", [False, True]) + @pytest.mark.parametrize( + "values,dtype", + [ + # uint8: true minimum is 0 (bottom of the unsigned range). + ([3, 1, 0, 5], torch.uint8), + # Signed dtypes at their exact minimum (INT_MIN), which overflows + # under the old ``x * -1`` negation trick. + ([3, 1, -128, 5], torch.int8), + ([3, 1, -(2**15), 5], torch.int16), + ([3, 1, -(2**31), 5], torch.int32), + # NOTE: coreai-torch narrows int64 inputs to int32 at the graph + # boundary, so the exercisable extreme for the int64 dtype *path* + # is int32's minimum (a true int64 min of -2**63 would be + # truncated at that boundary, unrelated to this argmin fix). + ([3, 1, -(2**31), 5], torch.int64), + ], + ) + async def test_dtype_min_value( + self, values: list[int], dtype: torch.dtype, keepdim: bool + ) -> None: + x = torch.tensor(values, dtype=dtype) + model = self.MinDimModel(dim=0, keepdim=keepdim).eval() + await validate_numerical_output(model=model, x=x) + + @pytest.mark.parametrize("keepdim", [False, True]) + @pytest.mark.parametrize( + "values,dtype", + [ + # Duplicate minima at the dtype extreme: torch returns the FIRST + # index (index 0 here), which the reversal must preserve. + ([-128, -128, 5], torch.int8), + ([0, 0, 5], torch.uint8), + ], + ) + async def test_duplicate_minima_first_index( + self, values: list[int], dtype: torch.dtype, keepdim: bool + ) -> None: + x = torch.tensor(values, dtype=dtype) + # Sanity-check the ground-truth tie-break we are asserting against. + assert int(torch.min(x, dim=0).indices) == 0 + model = self.MinDimModel(dim=0, keepdim=keepdim).eval() + await validate_numerical_output(model=model, x=x) + + @pytest.mark.parametrize("keepdim", [False, True]) + async def test_float_path_untouched(self, keepdim: bool) -> None: + # Float negation is exact/no-overflow; confirm it still validates. + x = torch.tensor([3.0, 1.0, -5.0, 2.0], dtype=torch.float32) + model = self.MinDimModel(dim=0, keepdim=keepdim).eval() + await validate_numerical_output(model=model, x=x) + + @pytest.mark.parametrize("dynamic", [False, True]) @pytest.mark.parametrize( "shape,dim,keepdim,dtype",