Repository: google/jax

Issue:
jnp.asarray does not copy when copy=True
This can be a problem now that we have donate_argnums. It took me a while to realize this, before I smartened up and just used x.copy() in the example below.

```
import functools
import jax
import jax.numpy as jnp

@functools.partial(jax.jit, donate_argnums=(0,))
def _test(array: jnp.ndarray):
  return array.at[0].set(77)

x = jnp.asarray([0, 1])
x_copy = jnp.array(x, copy=True)
new_x = _test(x)

# Gives: RuntimeError: Invalid argument: CopyToHostAsync() called on invalid buffer.
print(x_copy)

```

@tomhennigan 


Discussion:
@hawkinsp advises that we don't have in our device runtime a way to copy buffers! He'll add one, but in the meantime I'll add something that bounces back to the host.

Resolving patch:
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -43,7 +43,7 @@
 from ..abstract_arrays import UnshapedArray, ShapedArray, ConcreteArray, canonicalize_shape
 from ..config import flags
 from ..interpreters.xla import (DeviceArray, device_put, array_result_handler,
-                                DeviceValue)
+                                DeviceValue, abstractify)
 from ..interpreters.masking import Poly
 from .. import lax
 from .. import ops
@@ -2120,9 +2120,12 @@ def array(object, dtype=None, copy=True, order="K", ndmin=0):
     out = _device_put_raw(object)
     if dtype: assert _dtype(out) == dtype
   elif isinstance(object, (DeviceValue, core.Tracer)):
-    out = object
-    if dtype and _dtype(out) != dtype:
-      out = lax.convert_element_type(out, dtype)
+    if isinstance(object, DeviceArray) and copy:
+      # We perform a copy by bouncing back to the host
+      # TODO(phawkins): add a device runtime function to copy a buffer
+      out = _device_put_raw(np.asarray(object))
+    else:
+      out = object
   elif isinstance(object, (list, tuple)):
     if object:
       out = stack([array(elt, dtype=dtype) for elt in object])
@@ -2138,6 +2141,9 @@ def array(object, dtype=None, copy=True, order="K", ndmin=0):
 
     raise TypeError("Unexpected input type for array: {}".format(type(object)))
 
+  if dtype and _dtype(out) != dtype:
+    out = lax.convert_element_type(out, dtype)
+
   if ndmin > ndim(out):
     out = lax.broadcast(out, (1,) * (ndmin - ndim(out)))
   return out
@@ -2146,9 +2152,9 @@ def _can_call_numpy_array(x):
   return _all(not isinstance(l, (core.Tracer, DeviceValue))
               for l in tree_leaves(x))
 
+# TODO(mattjj): maybe move these two functions into xla.py
 def _device_put_raw(x):
-  aval = core.raise_to_shaped(core.get_aval(x))
-  return array_result_handler(None, aval)(device_put(x))
+  return array_result_handler(None, abstractify(x))(device_put(x))
 
 
 @_wraps(np.asarray)


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -3114,6 +3114,23 @@ def test_jit_nested_donate_ignored(self):
     with self.assertRaisesRegex(ValueError, "nested.*not supported"):
       jit_fun(a)
 
+  def test_jnp_array_copy(self):
+    # https://github.com/google/jax/issues/3412
+
+    @partial(api.jit, donate_argnums=(0,))
+    def _test(array):
+      return array.at[0].set(77)
+
+    x = jnp.asarray([0, 1])
+    x_copy = jnp.array(x, copy=True)
+    with warnings.catch_warnings():
+      warnings.simplefilter("ignore")
+      _test(x)  # donation
+
+    # Gives: RuntimeError: Invalid argument: CopyToHostAsync() called on invalid buffer.
+    print(x_copy)  # doesn't crash
+
+
   # === pmap ===
 
   @jtu.skip_on_devices("cpu", "gpu")  # In/out aliasing only supported on TPU.