Add support for device kwarg in astype to match Array API

Author: Micky774

jax/_src/numpy/lax_numpy.py

@@ -41,7 +41,7 @@
 import opt_einsum
 
 import jax
-from jax import jit
+from jax import jit, device_put
 from jax import errors
 from jax import lax
 from jax.sharding import Sharding, SingleDeviceSharding
@@ -2209,19 +2209,22 @@ def _convert_to_array_if_dtype_fails(x: ArrayLike) -> ArrayLike:
   else:
     return x
 
-
 @util.implements(getattr(np, "astype", None), lax_description="""
 This is implemented via :func:`jax.lax.convert_element_type`, which may
 have slightly different behavior than :func:`numpy.astype` in some cases.
 In particular, the details of float-to-int and int-to-float casts are
 implementation dependent.
 """)
-def astype(x: ArrayLike, dtype: DTypeLike | None, /, *, copy: bool = True) -> Array:
-  del copy  # unused in JAX
+def astype(x: ArrayLike, dtype: DTypeLike | None, /, *, copy: bool = True, device: xc.Device | Sharding | None = None) -> Array:
+  src_devices = x.devices() if hasattr(x, "devices") and not isinstance(x, core.Tracer) else None
+  if device is not None and src_devices != {device}:
+    return device_put(x, device)
+  arr = _array_copy(x) if copy else x
   if dtype is None:
     dtype = dtypes.canonicalize_dtype(float_)
   dtypes.check_user_dtype_supported(dtype, "astype")
-  return lax.convert_element_type(x, dtype)
+  return lax.convert_element_type(arr, dtype)
+
 
 
 @util.implements(np.asarray, lax_description=_ARRAY_DOC)

jax/experimental/array_api/_data_type_functions.py

@@ -13,12 +13,13 @@
 # limitations under the License.
 
 
+import builtins
 import functools
 from typing import NamedTuple
 import jax
+from jax._src.sharding import Sharding
 import jax.numpy as jnp
-
-
+from jax._src.lib import xla_client as xc
 from jax.experimental.array_api._dtypes import (
   bool, int8, int16, int32, int64, uint8, uint16, uint32, uint64,
   float32, float64, complex64, complex128
@@ -124,8 +125,8 @@ def _promote_types(t1, t2):
     raise ValueError("No promotion path for {t1} & {t2}")
 
 
-def astype(x, dtype, /, *, copy=True):
-  return jnp.array(x, dtype=dtype, copy=copy)
+def astype(x, dtype, /, *, copy: builtins.bool = True, device: xc.Device | Sharding | None = None):
+  return jnp.astype(x, dtype, copy=copy, device=device)
 
 
 def can_cast(from_, to, /):

jax/numpy/__init__.pyi

@@ -9,6 +9,8 @@ from jax._src import dtypes as _dtypes
 from jax._src.lax.lax import PrecisionLike
 from jax._src.lax.slicing import GatherScatterMode
 from jax._src.numpy.index_tricks import _Mgrid, _Ogrid, CClass as _CClass, RClass as _RClass
+from jax._src.sharding import Sharding
+from jax._src.lib import xla_client as xc
 from jax._src.typing import Array, ArrayLike, DType, DTypeLike, DimSize, DuckTypedArray, Shape
 from jax.numpy import fft as fft, linalg as linalg
 from jax.sharding import Sharding as _Sharding
@@ -112,7 +114,7 @@ def asarray(
 ) -> Array: ...
 def asin(x: ArrayLike, /) -> Array: ...
 def asinh(x: ArrayLike, /) -> Array: ...
-def astype(a: ArrayLike, dtype: Optional[DTypeLike], /, *, copy: builtins.bool = ...) -> Array: ...
+def astype(a: ArrayLike, dtype: Optional[DTypeLike], /, *, copy: builtins.bool = ..., device: xc.Device | Sharding | None = ...) -> Array: ...
 def atan(x: ArrayLike, /) -> Array: ...
 def atan2(x: ArrayLike, y: ArrayLike, /) -> Array: ...
 def atanh(x: ArrayLike, /) -> Array: ...