Repository: google/jax

Issue:
ConcreteArray has no attribute copy
I wanted to use `.copy()` on arrays to avoid working on views and corrupting values elsewhere.
During tracing jax called my function with a `ConcreteArray` and threw an error since these don't support `.copy()`.

Is there a different way i should have done this or should `ConcreteArray`s support a dummy `.copy()` that just returns `self`?


Discussion:
Can you share a full code example? How do you end up with a ConcreteArray?
@shoyer
A minimal reproducer:

```python
from jax import numpy as np, tree_map, tree_flatten, value_and_grad
from jax.numpy import ndarray as array
from typing import List


class PEPS:
    def __init__(self, tensors: List[array], *args):
        self._tensors = tree_map(lambda arr: arr.copy(), tensors)
        self.args = args

    def calculate_sth(self):
        return np.sum(tree_flatten(tree_map(np.linalg.norm, self._tensors))[0])


def cost_function(new_tensors):
    peps = PEPS(new_tensors)
    return peps.calculate_sth()


test_tensors = [np.zeros([i, i, i]) for i in range(5)]
val, grad = value_and_grad(cost_function)(test_tensors)
```

raises:
```
Traceback (most recent call last):
  File "/Users/jakobunfried/*CENSORED*/tmp/minimal_reproducer.py", line 21, in <module>
    val, grad = value_and_grad(cost_function)(test_tensors)
  File "/Users/jakobunfried/anaconda3/lib/python3.7/site-packages/jax/api.py", line 411, in value_and_grad_f
    ans, vjp_py = vjp(f_partial, *dyn_args)
  File "/Users/jakobunfried/anaconda3/lib/python3.7/site-packages/jax/api.py", line 1289, in vjp
    out_primal, out_vjp = ad.vjp(flat_fun, primals_flat)
  File "/Users/jakobunfried/anaconda3/lib/python3.7/site-packages/jax/interpreters/ad.py", line 109, in vjp
    out_primals, pvals, jaxpr, consts = linearize(traceable, *primals)
  File "/Users/jakobunfried/anaconda3/lib/python3.7/site-packages/jax/interpreters/ad.py", line 98, in linearize
    jaxpr, out_pvals, consts = pe.trace_to_jaxpr(jvpfun_flat, in_pvals)
  File "/Users/jakobunfried/anaconda3/lib/python3.7/site-packages/jax/interpreters/partial_eval.py", line 337, in trace_to_jaxpr
    jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
  File "/Users/jakobunfried/anaconda3/lib/python3.7/site-packages/jax/linear_util.py", line 152, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/Users/jakobunfried/*CENSORED*/tmp/minimal_reproducer.py", line 16, in cost_function
    peps = PEPS(new_tensors)
  File "/Users/jakobunfried/*CENSORED*/tmp/minimal_reproducer.py", line 8, in __init__
    self._tensors = tree_map(lambda arr: arr.copy(), tensors)
  File "/Users/jakobunfried/anaconda3/lib/python3.7/site-packages/jax/tree_util.py", line 124, in tree_map
    return treedef.unflatten(map(f, leaves))
  File "/Users/jakobunfried/*CENSORED*/tmp/minimal_reproducer.py", line 8, in <lambda>
    self._tensors = tree_map(lambda arr: arr.copy(), tensors)
  File "/Users/jakobunfried/anaconda3/lib/python3.7/site-packages/jax/core.py", line 365, in __getattr__
    attr = getattr(self.aval, name)
AttributeError: 'ConcreteArray' object has no attribute 'copy'

Process finished with exit code 1
```
@shoyer not sure if you get notified if I edit, so here goes another ping
Here's a simpler example:
```python
import jax
import numpy as np
jax.value_and_grad(lambda x: x.copy())(np.ones((2,)))
```

I guess one reason why `copy()` may not exist on `ConcreteArray` (or its base class `ShapedArray`) is that it the semantics of copying a tracer are not entirely clear:
1. should a copy return the same object on a tracer?
2. should a copy return a "disconnected copy"?

The later is not really possible for tracers, but it may be what users expect.
For context, it may be worth noting that JAX doesn't support in-place operations on arrays, so copies are generally not useful.

Nonetheless, it is maybe worth adding a dummy `copy()` method to facilitate duck typing. I think dask does this.
You are completely correct!
I was still in the `onp` mindset and trying to protect my class attributes from being a view of an array that could be in-place mutated elsewhere and thus change my class attribute arrays uncontrollably. But since jax doesn't do that, I don't even need to copy.
So for my use-case, just copying the list, i.e. `self._tensors = tensors.copy()` list should be completely fine.

Thanks!

In the end, it was useful for me to run into this error, since it brought my mistake to light.

For my part, I don't see that any changes are called for.
Do you want to close this, or is it still an open qustion if `ShapedArray` and its subclasses should support `.copy()` ?
> Do you want to close this, or is it still an open qustion if `ShapedArray` and its subclasses should support `.copy()` ?

I think this is still a valid open question.

Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -1383,11 +1383,10 @@ def device_put(x, device=None):
   return tree_map(lambda y: xla.device_put_p.bind(y, device=device), x)
 
 
-# TODO(mattjj): consider revising
 def _device_get(x):
   if isinstance(x, core.Tracer):
     return x
-  return x.copy()
+  return onp.asarray(x)
 
 def device_get(x):
   for y in tree_leaves(x):
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
@@ -1816,6 +1816,10 @@ def asarray(a, dtype=None, order=None):
   lax._check_user_dtype_supported(dtype, "asarray")
   return array(a, dtype=dtype, copy=False, order=order)
 
+@_wraps(onp.copy)
+def copy(a, order='K'):
+  return array(a, copy=True, order=order)
+
 
 @_wraps(onp.zeros_like)
 def zeros_like(x, dtype=None):
@@ -3436,7 +3440,8 @@ def _operator_round(number, ndigits=None):
 _diff_methods = ["clip", "compress", "conj", "conjugate", "cumprod", "cumsum",
                  "diagonal", "dot", "max", "mean", "min", "prod", "ptp",
                  "ravel", "repeat", "sort", "squeeze", "std", "sum",
-                 "swapaxes", "take", "tile", "trace", "transpose", "var"]
+                 "swapaxes", "take", "tile", "trace", "transpose", "var",
+                 "copy"]
 
 
 # Set up operator, method, and property forwarding on Tracer instances containing


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -1826,6 +1826,9 @@ def func1(x):
           re.DOTALL)):
       api.jit(func1)(2.)
 
+  def test_array_tracer_copy(self):
+    api.value_and_grad(lambda x: x.copy().sum())(np.ones(2))  # doesn't crash
+
 
 class JaxprTest(jtu.JaxTestCase):