code: def reshape_data_from_pmap(data: jnp.ndarray) -> jnp.ndarray:
      # data.shape = (#device, steps, #jobs/device, *)
      data = data.transpose([1, 0] + [i for i in range(2, data.ndim)])
      return jnp.reshape(
          data, (data.shape[0], data.shape[1] * data.shape[2], -1)
      )
  
imports: import logging
  from functools import partial
  from typing import Tuple
  from typing import Union
  import time
  import jax
  import jax.numpy
  from jax import random
  from jax.tree_util import tree_map
  from evojax.obs_norm import ObsNormalizer
  from evojax.task.base import TaskState
  from evojax.task.base import VectorizedTask
  from evojax.policy.base import PolicyState
  from evojax.policy.base import PolicyNetwork
  from evojax.util import create_logger
label_desc: Use image node_id 4 for calling a global function with example usage: reshape_data_from_pmap(data) and returns: jnp