code: def _to_device(
      self,
      device: xe.Device | Callable[[], xe.Device],
      /,
      *,
      stream: int | Any | None = None,
  ):
      if stream is not None:
          raise NotImplementedError(
              "stream argument of array.to_device()"
          )
      # The type of device is defined by Array.device. In JAX, this is a callable that
      # returns a device, so we must handle this case to satisfy the API spec.
      return jax.device_put(
          self, device() if callable(device) else device
      )
  
file_path: jax/jax/experimental/array_api/_array_methods.py
imports: from __future__ import annotations
  from typing import Any, Callable
  import jax
  from jax._src.array import ArrayImpl
  from jax.experimental.array_api._version import __array_api_version__
  from jax._src.lib import xla_extension
label_desc: Use image node_id 2 for calling a global function with example usage: _to_device() and returns: jax