code: def _pmap_device_order():
      return jax.local_devices()
  
imports: import collections
  import itertools
  import warnings
  from collections.abc import Iterable
  import jax
  import jax.numpy
  import numpy
  from jax import core, lax
  from jax.extend import linear_util
  from jax.interpreters import partial_eval
label_desc: Use image node_id 1 for calling a global function with example usage: _pmap_device_order() and returns: jax