Merge pull request #390 from ev-br/torch_repeat_short_ints

`torch.repeat` cast short ints to int32

Author: ev-br

array_api_compat/torch/_aliases.py

@@ -573,6 +573,10 @@ def count_nonzero(
 
 # "repeat" is torch.repeat_interleave;  also the dim argument
 def repeat(x: Array, repeats: int | Array, /, *, axis: int | None = None) -> Array:
+    if isinstance(repeats, torch.Tensor) and repeats.dtype in (torch.int8, torch.int16):
+        # torch rejects short integers for the `repeat` argument:
+        # https://github.com/pytorch/pytorch/issues/151311
+        repeats = repeats.to(torch.int32)
     return torch.repeat_interleave(x, repeats, axis)
 
 

numpy-1-22-xfails.txt

@@ -143,8 +143,6 @@ array_api_tests/test_operators_and_elementwise_functions.py::test_where_with_sca
 array_api_tests/test_has_names.py::test_has_names[creation-from_dlpack]
 array_api_tests/test_signatures.py::test_func_signature[from_dlpack]
 array_api_tests/test_signatures.py::test_array_method_signature[__dlpack__]
-# uint64 repeats not supported
-array_api_tests/test_manipulation_functions.py::test_repeat
 
 # 2024.12 support
 array_api_tests/test_signatures.py::test_func_signature[bitwise_and]

numpy-1-26-xfails.txt

@@ -39,8 +39,6 @@ array_api_tests/meta/test_hypothesis_helpers.py::test_symmetric_matrices
 # 2023.12 support
 array_api_tests/test_signatures.py::test_func_signature[from_dlpack]
 array_api_tests/test_signatures.py::test_array_method_signature[__dlpack__]
-# uint64 repeats not supported
-array_api_tests/test_manipulation_functions.py::test_repeat
 
 # 2024.12 support
 array_api_tests/test_data_type_functions.py::TestResultType::test_with_scalars

numpy-dev-xfails.txt

@@ -8,10 +8,6 @@ array_api_tests/test_data_type_functions.py::test_finfo[complex64]
 array_api_tests/test_signatures.py::test_func_signature[vecdot]
 array_api_tests/test_signatures.py::test_extension_func_signature[linalg.vecdot]
 
-# 2023.12 support
-# uint64 repeats not supported
-array_api_tests/test_manipulation_functions.py::test_repeat
-
 # Stubs have a comment: (**note**: libraries may return ``NaN`` to match Python behavior.); Apparently, NumPy does just that
 array_api_tests/test_special_cases.py::test_binary[floor_divide(x1_i is +infinity and isfinite(x2_i) and x2_i > 0) -> +infinity]
 array_api_tests/test_special_cases.py::test_binary[floor_divide(x1_i is +infinity and isfinite(x2_i) and x2_i < 0) -> -infinity]

numpy-xfails.txt

@@ -8,10 +8,6 @@ array_api_tests/test_data_type_functions.py::test_finfo[complex64]
 array_api_tests/test_signatures.py::test_func_signature[vecdot]
 array_api_tests/test_signatures.py::test_extension_func_signature[linalg.vecdot]
 
-# 2023.12 support
-# uint64 repeats not supported
-array_api_tests/test_manipulation_functions.py::test_repeat
-
 # Stubs have a comment: (**note**: libraries may return ``NaN`` to match Python behavior.); Apparently, NumPy does just that
 array_api_tests/test_special_cases.py::test_binary[floor_divide(x1_i is +infinity and isfinite(x2_i) and x2_i > 0) -> +infinity]
 array_api_tests/test_special_cases.py::test_binary[floor_divide(x1_i is +infinity and isfinite(x2_i) and x2_i < 0) -> -infinity]

torch-xfails.txt

@@ -139,8 +139,6 @@ array_api_tests/test_data_type_functions.py::test_finfo_dtype
 array_api_tests/test_data_type_functions.py::test_iinfo_dtype
 
 # 2023.12 support
-# https://github.com/pytorch/pytorch/issues/151311: torch.repeat_interleave rejects short integers
-array_api_tests/test_manipulation_functions.py::test_repeat
 # Argument 'device' missing from signature
 array_api_tests/test_signatures.py::test_func_signature[from_dlpack]
 # Argument 'max_version' missing from signature