code: def _adjust_classification_inputs(
      logits: jnp.ndarray, targets: jnp.ndarray, all_gather_loss: bool
  ) -> Tuple[jnp.ndarray, jnp.ndarray]:
      """Adjusts logits and label to have the same shape.
  
      Args:
        logits: Output of model in shape [batch, batch * max_num_captions,
          num_frames, 1].
        targets: Binary array of shape [batch, max_num_captions, num_frames, 1].
        all_gather_loss: Wether or not to gather results from all devices before
          computing metrics and loss.
  
      Returns:
        Reshaped logits and targets in shape (batch, batch, max_num_captions,
          num_frames, 1) when all_gather_loss = True. Otherwise, it is
          (batch, max_num_captions, num_frames, 1).
      """
      bs, num_cap, num_frames, _ = targets.shape
      reshaped_logits = logits.reshape((bs, bs, num_cap, num_frames, 1))
      if all_gather_loss:
          expanded_targets = jnp.zeros(
              (bs,) + targets.shape, dtype=targets.dtype
          )
          expanded_targets = expanded_targets.at[
              jnp.arange(bs), jnp.arange(bs)
          ].set(targets)
          return reshaped_logits, expanded_targets
      reshaped_logits = reshaped_logits[jnp.arange(bs), jnp.arange(bs)]
      return reshaped_logits, targets
  
imports: import functools
  from typing import Any, Dict, Mapping, Optional, Tuple, Union
  import immutabledict
  import jax.numpy
  import ml_collections
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import model_utils
  from scenic.projects.unloc import temporal_localization_base_model
label_desc: Use image node_id 1 for calling a global function with example usage: _adjust_classification_inputs(logits, targets, all_gather_loss) and returns: reshaped_logits,  targets, reshaped_logits,  expanded_targets