BUG: add torch.repeat

Author: ev-br

array_api_compat/torch/_aliases.py

@@ -177,7 +177,6 @@ def can_cast(from_: Union[DType, Array], to: DType, /) -> bool:
 # torch.conj sets the conjugation bit, which breaks conversion to other
 # libraries. See https://github.com/data-apis/array-api-compat/issues/173
 conj = torch.conj_physical
-
 # Two-arg elementwise functions
 # These require a wrapper to do the correct type promotion on 0-D tensors
 add = _two_arg(torch.add)
@@ -574,6 +573,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, 
@@ -854,6 +858,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

@@ -120,9 +120,9 @@ 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