code: def prepare_data(xs, pad=None, num_devices=None):
      """Makes sure data fits local devices.
  
      Args:
        xs: the data to be reshaped (tree-map'able).
        pad: if not None, zero-pad arrays such that their first dimension has this
          size. Also introduces a "mask" key in `xs` that contains ones where the
          un-padded data resides and zeros where the padding is.
        num_devices: Number of devices to distribute the data across. If None,
          ``local_device_count`` will be used instead.
  
      Returns:
        `xs` but re-shaped to (local_devices, -1, ...) and optionally padded.
      """
      # Create a mask which will be 1 for real entries and 0 for padded ones.
      if pad is not None:
          xs["mask"] = np.full(len(xs["image"]), 1.0)
      local_device_count = num_devices or jax.local_device_count()
  
      def _prepare(x):
          # Transforms x into read-only numpy array without copy if possible, see:
          # https://github.com/tensorflow/tensorflow/issues/33254#issuecomment-542379165
          x = np.asarray(memoryview(x))
  
          if pad is not None and len(x) != pad:
              # Append `pad - len(x)` rows of zeros.
              x = np.r_[
                  x, np.zeros((pad - len(x),) + x.shape[1:], x.dtype)
              ]
  
          return x.reshape((local_device_count, -1) + x.shape[1:])
  
      xs = jax.tree_util.tree_map(_prepare, xs)
      return {
          "inputs": xs["image"],
          "label": xs["label"],
          "batch_mask": xs["mask"],
      }
  
imports: import functools
  from typing import Any, Optional, Tuple
  from absl import logging
  from clu import metric_writers
  from flax import jax_utils
  import jax
  import jax.numpy
  import numpy
  from scenic.dataset_lib import dataset_utils
  from scenic.dataset_lib.big_transfer import builder
  from scenic.train_lib import train_utils
  import tensorflow_datasets
label_desc: Use image node_id 1 for calling a global function with example usage: prepare_data(xs, pad, num_devices) and returns: dict, x