Repository: google/jax

Issue:
`jax.numpy.array()` is much slower than `numpy.array()` when called on a list
As pointed out in the docs, converting lists to arrays is terribly slow compared to numpy:

```
import numpy as np
import jax.numpy as jnp

%%timeit
np.array([0] * int(1e6))
> 60.1 ms ± 578 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%%timeit
jnp.array([0] * int(1e6))
> 3min 12s ± 794 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
```

This is somewhat remedied by first calling numpy:
```
%%timeit
jnp.array(np.array([0] * int(1e6)))
> 61.9 ms ± 574 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
```

Ideally, one wouldn't want to make a proxy for `array()` by defining a function that always calls numpy first, so it would be nice to either explicitly do this on a low level, or improve the implementation of `array()` as to avoid this huge discrepancy in speed.


Discussion:
This is something worth fixing! 

@hawkinsp did a quick profile, I'll put the screenshot here just for posterity:

![image](https://user-images.githubusercontent.com/1458824/80822909-7ef8af80-8b90-11ea-8dc9-072e46c69eb6.png)

From the left half:

![image](https://user-images.githubusercontent.com/1458824/80823013-a8194000-8b90-11ea-8a4d-fccc3f777c41.png)

From the right half:

![image](https://user-images.githubusercontent.com/1458824/80822956-93d54300-8b90-11ea-80b4-32414e86695b.png)

Looks like a lot of dumb stuff going on. The good news is we should be able to improve it significantly :P
The same problem happens when converting from a bunch of `jnp.array`s to `np.array`. But the bulk of the time is in the creation of the `jnp.array`s.

```
import numpy as np
import jax.numpy as jnp

%%timeit
np.array([0] * int(1e3))
10000 loops, best of 3: 119 µs per loop

%%timeit
arr = jnp.array([0] * int(1e3))
1 loops, best of 3: 212 ms per loop

arr = jnp.array([0] * int(1e3))

%%timeit
np.array(arr)
100000 loops, best of 3: 7.1 µs per loop
```

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
@@ -35,20 +35,21 @@
 import numpy as np
 import opt_einsum
 
-from jax import jit, device_put, custom_jvp
+from jax import jit, custom_jvp
 from .vectorize import vectorize
 from ._util import _wraps
 from .. import core
 from .. import dtypes
 from ..abstract_arrays import UnshapedArray, ShapedArray, ConcreteArray, canonicalize_shape
 from ..config import flags
-from ..interpreters.xla import DeviceArray
+from ..interpreters.xla import (DeviceArray, device_put, array_result_handler,
+                                DeviceValue)
 from ..interpreters.masking import Poly
 from .. import lax
 from .. import ops
 from ..util import (partial, unzip2, prod as _prod,
                     subvals, safe_zip)
-from ..lib import pytree
+from ..tree_util import tree_leaves, tree_flatten
 
 FLAGS = flags.FLAGS
 flags.DEFINE_enum(
@@ -2107,19 +2108,24 @@ def array(object, dtype=None, copy=True, order="K", ndmin=0):
   if order is not None and order != "K":
     raise NotImplementedError("Only implemented for order='K'")
   lax._check_user_dtype_supported(dtype, "array")
-
-  if isinstance(object, ndarray) or isscalar(object):
-    out = device_put(object)
-    if dtype and _dtype(out) != dtypes.canonicalize_dtype(dtype):
+  dtype = dtype and dtypes.canonicalize_dtype(dtype)
+
+  if _can_call_numpy_array(object):
+    object = np.array(object, dtype=dtype, ndmin=ndmin)
+  assert type(object) not in dtypes.python_scalar_dtypes
+
+  if type(object) is np.ndarray:
+    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)
-  elif hasattr(object, '__array__'):
-    # this case is for duck-typed handling of objects that implement `__array__`
-    out = array(object.__array__(), dtype and dtypes.canonicalize_dtype(dtype))
   elif isinstance(object, (list, tuple)):
     if object:
       out = stack([array(elt, dtype=dtype) for elt in object])
     else:
-      out = array(np.array([], dtype or float_))
+      out = _device_put_raw(np.array([], dtype or float_))
   else:
     try:
       view = memoryview(object)
@@ -2134,6 +2140,15 @@ def array(object, dtype=None, copy=True, order="K", ndmin=0):
     out = lax.broadcast(out, (1,) * (ndmin - ndim(out)))
   return out
 
+def _can_call_numpy_array(x):
+  return _all(not isinstance(l, (core.Tracer, DeviceValue))
+              for l in tree_leaves(x))
+
+def _device_put_raw(x):
+  aval = core.raise_to_shaped(core.get_aval(x))
+  return array_result_handler(None, aval)(device_put(x))
+
+
 @_wraps(np.asarray)
 def asarray(a, dtype=None, order=None):
   lax._check_user_dtype_supported(dtype, "asarray")
@@ -3431,7 +3446,7 @@ def _split_index_for_jit(idx):
   # indexing logic to handle them.
   idx = _expand_bool_indices(idx)
 
-  leaves, treedef = pytree.flatten(idx)
+  leaves, treedef = tree_flatten(idx)
   dynamic = [None] * len(leaves)
   static = [None] * len(leaves)
   for i, x in enumerate(leaves):


Test patch:
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -2134,16 +2134,17 @@ def __array__(self, dtype=None):
             return np.array([], dtype=dtype)
     assert type(jnp.array(NDArrayLike())) == jax.interpreters.xla.DeviceArray
 
-    class DeviceArrayLike:
-        def __array__(self, dtype=None):
-            return jnp.array([], dtype=dtype)
-    assert type(jnp.array(DeviceArrayLike())) == jax.interpreters.xla.DeviceArray
+    # NOTE(mattjj): disabled b/c __array__ must produce ndarrays
+    # class DeviceArrayLike:
+    #     def __array__(self, dtype=None):
+    #         return jnp.array([], dtype=dtype)
+    # assert type(jnp.array(DeviceArrayLike())) == jax.interpreters.xla.DeviceArray
 
   def testArrayMethod(self):
     class arraylike(object):
       dtype = np.float32
       def __array__(self, dtype=None):
-        return 3.
+        return np.array(3., dtype=dtype)
     a = arraylike()
     ans = jnp.array(a)
     assert ans == 3.