BUG: torch: work around torch.round not supporting complex inputs

Author: ev-br

array_api_compat/torch/_aliases.py

@@ -912,6 +912,15 @@ def sign(x: Array, /) -> Array:
         return out
 
 
+def round(x: Array, /, **kwargs) -> Array:
+    # torch.round fails for complex inputs
+    # https://github.com/pytorch/pytorch/issues/58743#issuecomment-2727603845
+    if x.dtype.is_complex:
+        return torch.round(x.real, **kwargs) + 1j*torch.round(x.imag, **kwargs)
+    else:
+        return torch.round(x, **kwargs)
+
+
 def meshgrid(*arrays: Array, indexing: Literal['xy', 'ij'] = 'xy') -> tuple[Array, ...]:
     # torch <= 2.9 emits a UserWarning: "torch.meshgrid: in an upcoming release, it
     # will be required to pass the indexing argument."
@@ -923,7 +932,7 @@ def meshgrid(*arrays: Array, indexing: Literal['xy', 'ij'] = 'xy') -> tuple[Arra
            'permute_dims', 'bitwise_invert', 'newaxis', 'conj', 'add',
            'atan2', 'bitwise_and', 'bitwise_left_shift', 'bitwise_or',
            'bitwise_right_shift', 'bitwise_xor', 'copysign', 'count_nonzero',
-           'diff', 'divide',
+           'diff', 'divide', 'round',
            'equal', 'floor_divide', 'greater', 'greater_equal', 'hypot',
            'less', 'less_equal', 'logaddexp', 'maximum', 'minimum',
            'multiply', 'not_equal', 'pow', 'remainder', 'subtract', 'max',

torch-xfails.txt

@@ -130,7 +130,6 @@ array_api_tests/test_statistical_functions.py::test_var
 
 
 # These functions do not yet support complex numbers
-array_api_tests/test_operators_and_elementwise_functions.py::test_round
 array_api_tests/test_set_functions.py::test_unique_counts
 array_api_tests/test_set_functions.py::test_unique_values