code: def shard(pytree, n_devices=None):
      """Reshapes all arrays in the pytree to add a leading n_devices dimension.
  
      To be used for pmap-based data-parallelism.
  
      Note: We assume that all arrays in the pytree have leading dimension divisible
      by n_devices and reshape (host_batch_size, height, width, channel) to
      (local_devices, device_batch_size, height, width, channel).
  
      Args:
        pytree: A pytree of arrays to be sharded.
        n_devices: If None, this will be set to jax.local_device_count().
  
      Returns:
        Sharded data.
      """
      if n_devices is None:
          n_devices = jax.local_device_count()
  
      def _shard_array(array):
          return array.reshape((n_devices, -1) + array.shape[1:])
  
      return jax.tree_util.tree_map(_shard_array, pytree)
  
imports: import collections
  import dataclasses
  import functools
  import itertools
  from typing import Any, Callable, Dict, Iterator, Optional, Sequence, Union
  from absl import logging
  from flax.training import common_utils
  import jax
  import jax.numpy
  import numpy
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 2 for calling a global function with example usage: shard(pytree, n_devices) and returns: jax, array