Add support for copy kwarg in astype to match Array API

Author: Micky774

jax/_src/numpy/lax_numpy.py

@@ -2270,28 +2270,25 @@ def astype(x: ArrayLike, dtype: DTypeLike | None,
   if dtype is None:
     dtype = dtypes.canonicalize_dtype(float_)
   dtypes.check_user_dtype_supported(dtype, "astype")
-  if (
-    issubdtype(x_arr.dtype, complexfloating)
-    and dtypes.isdtype(dtype, ("integral", "real floating"))
-  ):
-    warnings.warn(
-      "Casting from complex to real dtypes will soon raise a ValueError. "
-      "Please first use jnp.real or jnp.imag to take the real/imaginary "
-      "component of your input.",
-      DeprecationWarning, stacklevel=2
-    )
-
-  out = x_arr
-  # convert_element_type(complex, bool) has the wrong semantics.
-  if np.dtype(dtype) == bool and issubdtype(x_arr.dtype, complexfloating):
-    out = (x_arr != _lax_const(x_arr, 0))
-  out = _place_array(out, device=device, copy=copy)
+  if issubdtype(x_arr.dtype, complexfloating):
+    if dtypes.isdtype(dtype, ("integral", "real floating")):
+      warnings.warn(
+        "Casting from complex to real dtypes will soon raise a ValueError. "
+        "Please first use jnp.real or jnp.imag to take the real/imaginary "
+        "component of your input.",
+        DeprecationWarning, stacklevel=2
+      )
+    elif np.dtype(dtype) == bool:
+      # convert_element_type(complex, bool) has the wrong semantics.
+      x_arr = (x_arr != _lax_const(x_arr, 0))
+
+  x_arr = _place_array(x_arr, device=device, copy=copy)
 
   # We offer a more specific warning than the usual ComplexWarning so we prefer
   # to issue our warning.
   with warnings.catch_warnings():
     warnings.simplefilter("ignore", ComplexWarning)
-    return lax.convert_element_type(out, dtype)
+    return lax.convert_element_type(x_arr, dtype)
 
 def _place_array(x, device=None, copy=None):
   # TODO(micky774): Implement in future PRs as we formalize device placement