code: def maybe_pad_batch(
      batch: Dict[str, PyTree],
      train: bool,
      batch_size: int,
      pixel_level: bool = False,
      inputs_key: str = "inputs",
      batch_dim: int = 0,
  ) -> Dict[str, jnp.ndarray]:
      """Zero pad the batch on the right to the batch_size.
  
      All leave tensors in the batch pytree will be padded. This function expects
      the root structure of the batch pytree to be a dictionary and returns a
      dictionary with the same structure (and substructures), additionally with the
      key 'batch_mask' added to the root dict, with 1.0 indicating indices which are
      true data and 0.0 indicating a padded index. `batch_mask` will be used for
      calculating the weighted cross entropy, or weighted accuracy.
  
      Note that in this codebase, we assume we drop the last partial batch from the
      training set, so if the batch is from the training set (i.e. `train=True`),
      or when the batch is from the test/validation set, but it is a complete batch,
      we *modify* the batch dict by adding an array of ones as the `batch_mask` of
      all examples in the batch. Otherwise, we create a new dict that has the padded
      patch and its corresponding `batch_mask` array.
  
      Note that batch_mask can be also used as the label mask (not input mask), for
      task that are pixel/token level. This is simply done by applying the mask we
      make for padding the partial batches on top of the existing label mask.
  
      Args:
        batch: A dictionary containing a pytree. If `inputs_key` is not set, we use
          the first leave to get the current batch size. Otherwise, the tensor
          mapped with `inputs_key` at the root dictionary is used.
        train: if the batch is from the training data. In that case, we drop
          the last (incomplete) batch and thus don't do any padding.
        batch_size: All arrays in the dict will be padded to have first
          dimension equal to desired_batch_size.
        pixel_level: If True, this will create a pixel-level (instead of
          example-level) mask, e.g. for segmentation models.
        inputs_key: Indicating the key used for the input that we do batch padding
          based on.
        batch_dim: Batch dimension. The default is 0, but it can be different
          if a sharded batch is given.
  
      Returns:
        A dictionary mapping the same keys to the padded batches. Additionally, we
        add a key representing weights, to indicate how the batch was padded.
      """
      assert (
          batch_dim >= 0
      ), f"batch_dim=={batch_dim} is expected to be >= 0"
      if inputs_key is None:
          sample_tensor = jax.tree_util.tree_leaves(batch)[0]
      else:
          sample_tensor = batch[inputs_key]
      if sample_tensor.shape[batch_dim] > batch_size:
          raise ValueError(
              f"The indicated target batch_size is {batch_size}, but "
              "the size of the current batch is larger than that: "
              f"{sample_tensor.shape[batch_dim]}."
          )
      batch_pad = batch_size - sample_tensor.shape[batch_dim]
  
      if pixel_level:
          unpadded_mask_shape = sample_tensor.shape[:-1]
      else:
          assert "batch_mask" not in batch, (
              "When the labels of the task are not pixel-level, batch_mask should "
              "not be already present in the batch."
          )
          unpadded_mask_shape = sample_tensor.shape[: batch_dim + 1]
  
      if train and batch_pad != 0:
          raise ValueError(
              "In this codebase, we assumed that we always drop the "
              "last partial batch of the train set. Please use "
              "` drop_remainder=True` for the training set."
          )
  
      # Most batches will not need padding, so we quickly return to avoid slowdown.
      if train or batch_pad == 0:
          if "batch_mask" not in batch:
              batch["batch_mask"] = np.ones(
                  unpadded_mask_shape, dtype=np.float32
              )
          return batch
  
      def zero_pad(array):
          pad_with = (
              [(0, 0)] * batch_dim
              + [(0, batch_pad)]
              + [(0, 0)] * (array.ndim - batch_dim - 1)
          )
          return np.pad(array, pad_with, mode="constant")
  
      padded_batch = jax.tree_util.tree_map(zero_pad, batch)
      padded_batch_mask = zero_pad(
          np.ones(unpadded_mask_shape, dtype=np.float32)
      )
      if "batch_mask" in padded_batch:
          padded_batch["batch_mask"] *= padded_batch_mask
      else:
          padded_batch["batch_mask"] = padded_batch_mask
      return padded_batch
  
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 1 for calling a global function with example usage: maybe_pad_batch(batch, train, batch_size, pixel_level, inputs_key, batch_dim) and returns: padded_batch, batch, np