Merge pull request #305 from ev-br/torch_repeat

BUG: add torch.repeat

Author: ev-br

array_api_compat/torch/_aliases.py

@@ -555,6 +555,11 @@ def count_nonzero(
         return result
 
 
+# "repeat" is torch.repeat_interleave;  also the dim argument
+def repeat(x: Array, repeats: int | Array, /, *, axis: int | None = None) -> Array:
+    return torch.repeat_interleave(x, repeats, axis)
+
+
 def where(
     condition: Array, 
     x1: Array | bool | int | float | complex, 
@@ -835,6 +840,6 @@ def sign(x: Array, /) -> Array:
            'UniqueAllResult', 'UniqueCountsResult', 'UniqueInverseResult',
            'unique_all', 'unique_counts', 'unique_inverse', 'unique_values',
            'matmul', 'matrix_transpose', 'vecdot', 'tensordot', 'isdtype',
-           'take', 'take_along_axis', 'sign', 'finfo', 'iinfo']
+           'take', 'take_along_axis', 'sign', 'finfo', 'iinfo', 'repeat']
 
 _all_ignore = ['torch', 'get_xp']

torch-xfails.txt

@@ -124,9 +124,8 @@ 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
-array_api_tests/test_has_names.py::test_has_names[manipulation-repeat]
+# https://github.com/pytorch/pytorch/issues/151311: torch.repeat_interleave rejects short integers
 array_api_tests/test_manipulation_functions.py::test_repeat
-array_api_tests/test_signatures.py::test_func_signature[repeat]
 # Argument 'device' missing from signature
 array_api_tests/test_signatures.py::test_func_signature[from_dlpack]
 # Argument 'max_version' missing from signature