code: def reshape_by_device(x, n_devices, pure_np=False):
      """Reshapes possibly nested ``x`` into a shape ``(n_devices, ...)``."""
  
      def f(x):
          x_shape = list(x.shape)
          batch_size = x_shape[0]
          batch_size_per_device = batch_size // n_devices
          if batch_size_per_device * n_devices != batch_size:
              raise ValueError(
                  f"Number of devices ({n_devices}) does not evenly "
                  f"divide batch size ({batch_size})."
              )
          new_shape_prefix = [n_devices, batch_size_per_device]
          if pure_np:
              return np.reshape(x, new_shape_prefix + x_shape[1:])
          else:
              return jnp.reshape(x, new_shape_prefix + x_shape[1:])
  
      return fastmath.nested_map(f, x)
  
imports: import jax
  import numpy
  from trax import fastmath
  from trax.fastmath import numpy
  from trax.layers import base
label_desc: Use image node_id 5 for calling a global function with example usage: reshape_by_device(x, n_devices, pure_np) and returns: fastmath, np, jnp