code: def split_shape_and_device(array):
      if (
          isinstance(array, jax.Array)
          and hasattr(array, "sharding")
          and isinstance(array.sharding, jax.sharding.PmapSharding)
      ):
          return array.shape[0], array.shape[1:]
      else:
          return None, array.shape
  
label_desc: Use image node_id 1 for calling a global function with example usage: split_shape_and_device(array) and returns: array,  array, None,  array