Raise on implicit device transfer in `__setitem__`

[ogrisel]

Consider the following two arrays:

>>> import array_api_strict as xp
>>> a = xp.ones(10, device=xp.Device("device2"))
>>> b = 2 * xp.ones(10)

As expected, it's not possible to run operations that combine two such arrays because they live on different devices:

>>> a + b
Traceback (most recent call last):
  Cell In[23], line 1
    a + b
  File ~/miniforge3/envs/dev/lib/python3.13/site-packages/array_api_strict/_array_object.py:545 in __add__
    self._check_type_device(other)
  File ~/miniforge3/envs/dev/lib/python3.13/site-packages/array_api_strict/_array_object.py:241 in _check_type_device
    raise ValueError(f"Arrays from two different devices ({self.device} and {other.device}) can not be combined.")
ValueError: Arrays from two different devices (array_api_strict.Device('device2') and array_api_strict.Device('CPU_DEVICE')) can not be combined.

However, doing assignment works without raising an exception despite the fact that this causes an implicit data transfer between devices:

>>> b[:5] = a[:5]
>>> b
Array([1., 1., 1., 1., 1., 2., 2., 2., 2.,
       2.], dtype=array_api_strict.float64)

Shouldn't the __setitem__ call raise in this case?

[ogrisel]

Note that the current behavior (no exception on cross-device __setitem__) of array-api-strict is aligned with PyTorch:

>>> import torch
>>> a = torch.ones(10, device="mps")
>>> b = 2 * torch.ones(10)  # CPU
>>> a + b
Traceback (most recent call last):
  Cell In[9], line 1
    a + b
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, mps:0 and cpu!
>>> b[:5] = a[:5]  # implicit MPS to CPU transfer does not raise!
>>> b
tensor([1., 1., 1., 1., 1., 2., 2., 2., 2., 2.])

Not sure if this lack of exception could be considered a "bug" of PyTorch with MPS.

However, DPNP on some hardware architectures do raise: ValueError: Execution queue is not compatible with allocation queues as seen in scikit-learn/scikit-learn#32460 (comment).

[ogrisel]

I asked Claude that did a source code / doc review, and __setitem__ is considered "data movement" operation (like .to(), .copy_(), .cpu(), .cuda(), ...) and therefore is not expected to raise in PyTorch (contrary to "compute" operations like +).

I also checked the array API spec and there is no mention of constraints related to device transfers: https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__setitem__.html#setitem.

So I am not sure if it's ok or not for an array library to raise in this case.

[ogrisel]

Also from:

• https://data-apis.org/array-api/2025.12/design_topics/device_support.html

Handling devices is complex, and some frameworks have elaborate policies for handling device placement. Therefore this section only gives recommendations, rather than hard requirements:
...

• Raise an exception if an operation involves arrays on different devices (i.e. avoid implicit data transfer between devices).

So I think it would make sense for array-api-strict to raise on cross-device __setitem__ if we consider assignment an "operation".

[ogrisel]

cc @betatim.