diff --git a/src/array_api_compat/torch/_aliases.py b/src/array_api_compat/torch/_aliases.py index e27c3ca2..af6b98a5 100644 --- a/src/array_api_compat/torch/_aliases.py +++ b/src/array_api_compat/torch/_aliases.py @@ -585,15 +585,19 @@ def where(condition: Array, x1: Array | complex, x2: Array | complex, /) -> Arra return torch.where(condition, x1, x2) -# torch.reshape doesn't have the copy keyword def reshape(x: Array, /, shape: tuple[int, ...], *, copy: bool | None = None, **kwargs: object) -> Array: - if copy is not None: - raise NotImplementedError("torch.reshape doesn't yet support the copy keyword") + if copy is True: + return x.clone(memory_format=torch.contiguous_format).view(shape, **kwargs) + if copy is False: + try: + return x.view(shape, **kwargs) + except RuntimeError as error: + raise ValueError(str(error)) from error return torch.reshape(x, shape, **kwargs) # torch.arange doesn't support returning empty arrays diff --git a/tests/test_torch.py b/tests/test_torch.py index 3d6ebc46..a532ab4c 100644 --- a/tests/test_torch.py +++ b/tests/test_torch.py @@ -12,6 +12,21 @@ from array_api_compat import torch as xp +def test_reshape_copy(): + array = xp.arange(12).reshape((3, 4)) + + view = xp.reshape(array, (4, 3), copy=False) + view[0, 0] = -1 + assert array[0, 0] == -1 + + with pytest.raises(ValueError): + xp.reshape(array.T, (12,), copy=False) + + copied = xp.reshape(array, (4, 3), copy=True) + copied[0, 0] = -2 + assert array[0, 0] == -1 + + class TestResultType: def test_empty(self): with pytest.raises(ValueError):