code: def as_native_dev(device, /):
      if not isinstance(device, str):
          return device
      dev_split = ivy.Device(device).split(":")
      device = dev_split[0]
      if len(dev_split) > 1:
          idx = int(dev_split[1])
      else:
          idx = 0
      return jax.devices(device)[idx]
  
imports: import os
  import jax
  from typing import Union, Optional
  import jaxlib.xla_extension
  import ivy
  from ivy.functional.backends.jax import JaxArray
  from ivy.functional.ivy.device import _shift_native_arrays_on_default_device, Profiler
label_desc: Use image node_id 6 for calling a global function with example usage: as_native_dev() and returns: jax, device