code: def shard_jit(data: PyTree, global_devices: np.ndarray) -> PyTree:
      """Shards data for use in jit-based pipelines.
  
      Note that the order of global devices for sharding data is important and
      should be compatible with device order used in the rest of the trainer for
      models params, state, etc.
  
      Based on:
      https://github.com/google-research/big_vision/blob/main/big_vision/input_pipeline.py.
  
      Args:
        data: PyTree of data
        global_devices: List of global devices to shard over.
  
      Returns:
        Sharded data.
      """
  
      def _shard_array(x):
          mesh = jax.sharding.Mesh(global_devices, ("devices",))
          sharding = jax.sharding.NamedSharding(
              mesh, jax.sharding.PartitionSpec("devices")
          )
          local_ds = mesh.local_devices
  
          x = np.asarray(
              memoryview(x)
          )  # No-copy: http://shortn/_KM5whIEtWI
          xs = jax.device_put(
              np.split(x, len(local_ds), axis=0), local_ds
          )
  
          global_shape = (
              x.shape[0] * jax.process_count(),
              *x.shape[1:],
          )
          return jax.make_array_from_single_device_arrays(
              global_shape, sharding, xs
          )
  
      return jax.tree_util.tree_map(_shard_array, data)
  
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 3 for calling a global function with example usage: shard_jit(data, global_devices) and returns: jax, jax