Expand support for __jax_array__ in jnp.array.

This relates to the long discussion in google#4725 and google#10065.

jax/_src/numpy/lax_numpy.py

@@ -1862,6 +1862,8 @@ def array(object: Any, dtype: Optional[DTypeLike] = None, copy: bool = True,
   if isinstance(object, (bool, int, float, complex)):
     _ = dtypes.coerce_to_array(object, dtype)
 
+  object = tree_map(lambda leaf: leaf.__jax_array__() if hasattr(leaf, "__jax_array__") else leaf,
+                    object)
   leaves = tree_leaves(object)
   if dtype is None:
     # Use lattice_result_type rather than result_type to avoid canonicalization.

tests/api_test.py

@@ -3574,6 +3574,13 @@ def __jax_array__(self):
     for f in [jnp.isscalar, jnp.size, jnp.shape, jnp.dtype]:
       self.assertEqual(f(x), f(a))
 
+    x = AlexArray(jnp.array(1))
+    a1 = jnp.array(x)
+    self.assertAllClose(1, a1)
+
+    a2 = jnp.array(((x, x), [x, x]))
+    self.assertAllClose(np.array(((1, 1), (1, 1))), a2)
+
   def test_constant_handler_mro(self):
     # https://github.com/google/jax/issues/6129