code: def _get_device_by_id(device_id: int) -> xc.Device:
      for device in jax.devices():
          if device.id == device_id:
              return device
      raise ValueError(f"Device {device_id} was not found")
  
imports: import pickle
  import unittest
  from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  from jax import numpy
  from jax import config
  from jax.interpreters import pxla
  from jax._src import test_util
  from jax._src.lib import xla_client
  import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: _get_device_by_id(device_id) and returns: device