code: def _adjust_regression_inputs(
      displacements: jnp.ndarray,
  ) -> jnp.ndarray:
      """Adjusts prediction to have the same shape with ground truth.
  
      Regression loss is only computed at positive frames.
  
      Args:
        displacements: Predicted displacements in shape [batch, batch *
          max_num_captions, num_frames, 2].
  
      Returns:
        Reshaped prediction in shape (batch, max_num_captions, num_frames, 2)
      """
      bs, num_caps, num_frames, _ = displacements.shape
      num_caps_per_vid = num_caps // bs
      reshaped_displacements = displacements.reshape(
          (bs, bs, num_caps_per_vid, num_frames, 2)
      )
      return reshaped_displacements[jnp.arange(bs), jnp.arange(bs)]
  
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 2 for calling a global function with example usage: _adjust_regression_inputs(displacements) and returns: reshaped_displacements