Repository: google/jax

Issue:
jax.numpy.array should allow specifying placement properties (or docs should say why not)
A jax.numpy.array has a device and a committed/uncommitted state, so it seems strange that jax.numpy.array() does not have a way to specify these aspects of jax array state when an array is created.

The specific use case I have is that I want to convert a very large regular numpy array into a jax.numpy.array placed on the CPU device, ideally without copying it.

So it seems that I ought to be able to do something like:
jax_array = jax.numpy.array(numpy_array, copy=False, device=jax.devices("cpu")[0])

Thus, this feature request is to provide new parameters for jax.numpy.array to allow the device and commited/uncommited state to be specified.

If it's impossible or difficult to provide this functionality, the jax.numpy.array documentation ought to be amended to discuss this point, since it is a "surprise".  Ideally the docs would also suggest the appropriate idiom(s) for handling this case, particularly in Eager mode.

I think I can find other workarounds for this problem in my particular case, so it's not a blocking issue, but I'm filing the bug because this situation feels like bad API ergonomics to me.


Discussion:
Thanks for the report. I think the issue here is that the bare-metal way to create & control `DeviceArray` properties is using [`jax.device_put`](https://jax.readthedocs.io/en/latest/jax.html#jax.device_put), whereas `jnp.array` and similar functions are just wrappers of the NumPy API for convenience and familiarity. `np.array` doesn't have any `device` argument so `jnp.array` does not either.

Given that `jnp.array` has become by far the most used API for creating JAX arrays, we could talk about extending the API in this case rather than making people use the less familiar `device_put` API. What do you think? If `device_put` were mentioned in the `jnp.array` docstring, would that be sufficient to solve the issue you encountered?
Either fix would be acceptable.  Looking just at this specific issue, I would slightly prefer that the jnp.array() call be extended, but I don't know what your global policy is on "extending" the np.XXX calls. If you consistently avoid extending those calls, then I wouldn't advocate breaking consistency for this case, since there is an acceptable doc fix.
One more thought:  I understand that it's possible to use a jnp.array followed by a jax.device_put.  But in eager mode, won't that lead to the array being temporarily stored on the "wrong" device?  So perhaps it is necessary to extend jnp.array() to avoid this problem.
The recommended approach would be to avoid `jnp.array` and use `device_put` directly (assuming your input is a numpy array). If you don't have a numpy array, you can first convert your input to a numpy array before passing it to `device_put`. This is roughly what `jnp.array` does.
In case it's helpful, you can find more information on device placement in the FAQ: https://jax.readthedocs.io/en/latest/faq.html#controlling-data-and-computation-placement-on-devices

Resolving patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -3515,8 +3515,16 @@ def atleast_3d(*arys):
     return [atleast_3d(arr) for arr in arys]
 
 
-@_wraps(np.array)
-def array(object, dtype=None, copy=True, order="K", ndmin=0):
+_ARRAY_DOC = """
+This function will create arrays on JAX's default device. For control of the
+device placement of data, see :func:`jax.device_put`. More information is
+available in the JAX FAQ at :ref:`faq-data-placement` (full FAQ at
+https://jax.readthedocs.io/en/latest/faq.html).
+"""
+
+
+@_wraps(np.array, lax_description=_ARRAY_DOC)
+def array(object, dtype=None, copy=True, order="K", ndmin=0, *, device=None):
   if order is not None and order != "K":
     raise NotImplementedError("Only implemented for order='K'")
 
@@ -3577,7 +3585,7 @@ def _can_call_numpy_array(x):
               for l in tree_leaves(x))
 
 
-@_wraps(np.asarray)
+@_wraps(np.asarray, lax_description=_ARRAY_DOC)
 def asarray(a, dtype=None, order=None):
   lax._check_user_dtype_supported(dtype, "asarray")
   dtype = dtypes.canonicalize_dtype(dtype) if dtype is not None else dtype


Test patch: