Repository: google/jax

Issue:
Ops involving only CPU tensors run on device, not CPU
In the code below running jax op that has CPU backed inputs runs on device, (which for my use case OOMs):
```
def CpuArray(numpy_array):
  return jax.device_put(numpy_array, device=jax.devices(backend='cpu')[0])

data = CpuArray(np.ones([100, 4], dtype=np.float32))
slice_indexes = CpuArray(np.zeros([100], dtype=np.int32))
# Summing the data runs on TPU.
summed_data = jnp.sum(data[slice_indexes]) 
# So does slicing it.
sliced_data = data[100]

print("data device", data.device_buffer.device())
print("sliced indexes", slice_indexes.device_buffer.device())
print("sliced data", sliced_data.device_buffer.device())
print("summed data", summed_data.device_buffer.device())
```

Prints:

```
data device cpu:0
sliced indexes device cpu:0
Both output arrays are on TPU:
sliced data device TPU_0(host=0,(0,0,0,0))
summed data device TPU_0(host=0,(0,0,0,0))
```



Discussion:
This feels similar to #2883
I agree! I'm not sure yet if it's the same or not.
I verified that on master on a GPU machine this prints gpu:0:

```python
import jax
import numpy as onp

CPU = jax.devices('cpu')

x = jax.device_put(onp.ones(2), CPU[0])
y = x[0]
print(y.device_buffer.device())
```
Possibly also related to #2878.
I think it's because [`_rewriting_take` calls `asarray`](https://github.com/google/jax/blob/3216f5ca4647a00a5a53663846a56dc0faf1fc01/jax/numpy/lax_numpy.py#L3190), which is itself not respecting the device of its input.

Resolving patch:
diff --git a/jax/interpreters/xla.py b/jax/interpreters/xla.py
--- a/jax/interpreters/xla.py
+++ b/jax/interpreters/xla.py
@@ -16,7 +16,7 @@
 from collections import defaultdict
 import itertools as it
 import operator as op
-from typing import Any, Callable, Dict, Sequence, Type
+from typing import Any, Callable, Dict, Sequence, Type, Optional
 
 from absl import logging
 import numpy as onp
@@ -920,16 +920,23 @@ def _device_put_device_array(x, device):
   return _force(x).device_buffer
 device_put_handlers[DeviceArray] = _device_put_device_array
 
-def _copy_device_array_to_device(x, device):
-  if is_device_constant(x):
+def _copy_device_array_to_device(x: DeviceArray, device: Optional[xc.Device]):
+  if device is None:
+    # no copying to be done because there's no target specified
+    return x
+  elif is_device_constant(x):
+    # create a new DeviceArray with the same lazy expr, no copying
     return DeviceArray(x.aval, device, x._lazy_expr, DeviceConstant(device))
   elif xb.get_device_backend(device).platform == x.device_buffer.platform():
-    if device is None or x.device_buffer.device() == device:
+    # source and target platforms are the same
+    if x.device_buffer.device() == device:
+      # no copying to be done because source equals target
       return x
     else:
+      # move the buffer with a device-to-device copy
       moved_buf = x.device_buffer.copy_to_device(device)
   else:
-    # Buffers from different XLA backends are passed through the host.
+    # buffers from different XLA backends are passed through the host.
     backend = xb.get_device_backend(device)
     moved_buf = backend.buffer_from_pyval(x.device_buffer.to_py(), device)
   return DeviceArray(x.aval, device, x._lazy_expr, moved_buf)


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -323,9 +323,6 @@ def test_device_put_across_platforms(self):
       x = api.device_put(val, device=cpu_device)
       self.assertEqual(x.device_buffer.device(), cpu_device)
 
-    y = api.device_put(x)
-    self.assertEqual(y.device_buffer.device(), default_device)
-
   def test_jit_on_all_devices(self):
     # Verifies we can run the same computation on every device present, even
     # if they are, for example, different models of GPU.
diff --git a/tests/multibackend_test.py b/tests/multibackend_test.py
--- a/tests/multibackend_test.py
+++ b/tests/multibackend_test.py
@@ -171,6 +171,24 @@ def my_sin(x): return np.sin(x)
     result4 = api.jit(my_sin, backend="cpu")(2)
     self.assertEqual(result4.device_buffer.device(), cpus[0])
 
+  @jtu.skip_on_devices("cpu")  # test only makes sense on non-cpu backends
+  def test_indexing(self):
+    # https://github.com/google/jax/issues/2905
+    cpus = api.devices("cpu")
+
+    x = api.device_put(onp.ones(2), cpus[0])
+    y = x[0]
+    self.assertEqual(y.device_buffer.device(), cpus[0])
+
+  @jtu.skip_on_devices("cpu")  # test only makes sense on non-cpu backends
+  def test_sum(self):
+    # https://github.com/google/jax/issues/2905
+    cpus = api.devices("cpu")
+
+    x = api.device_put(onp.ones(2), cpus[0])
+    y = x.sum()
+    self.assertEqual(y.device_buffer.device(), cpus[0])
+
 
 if __name__ == "__main__":
   absltest.main()