Add support for device/copy kw in from_dlpack to match Array API

Author: Micky774

jax/numpy/__init__.pyi

@@ -12,7 +12,6 @@ from jax._src.numpy.index_tricks import _Mgrid, _Ogrid, CClass as _CClass, RClas
 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
-import jaxlib.xla_client as xc
 import numpy as _np
 
 _T = TypeVar('_T')
@@ -354,7 +353,7 @@ def fmax(x: ArrayLike, y: ArrayLike, /) -> Array: ...
 def fmin(x: ArrayLike, y: ArrayLike, /) -> Array: ...
 def fmod(x: ArrayLike, y: ArrayLike, /) -> Array: ...
 def frexp(x: ArrayLike, /) -> tuple[Array, Array]: ...
-def from_dlpack(x: Any, /, *, device: _Sharding | xc.Device | None = None, copy: builtins.bool | None = None) -> Array: ...
+def from_dlpack(x: Any, /, *, device: _Sharding | _Device = Any | None = None, copy: builtins.bool | None = None) -> Array: ...
 def frombuffer(buffer: Union[bytes, Any], dtype: DTypeLike = ...,
                count: int = ..., offset: int = ...) -> Array: ...
 def fromfile(*args, **kwargs): ...