code: def move_true_to_pred_mask(*, true_mask, true_box, pred_box):
      """Scales and translates the true mask into the pred box reference frame.
  
      If predicted and true boxes are not identical, then the segmentation masks
      within the boxes will be misaligned. This function scales and translates the
      true mask so that it aligns with the predicted mask, based on the positions
      of the true and predicted bounding boxes.
  
      This decouples the mask loss from the bounding box locations. In other words,
      the mask loss will be exactly as if we were using full-image masks, rather
      than within-box masks.
  
      Args:
        true_mask: Array of shape (n, n) containing the true mask for the area
          inside the true bounding box.
        true_box: True (cx, cy, w, h) bounding box.
        pred_box: Predicted (cx, cy, w, h) bounding box.
  
      Returns:
        New array of shape (n, n) containing the true mask for the area inside the
        predicted box.
  
      Raises:
        ValueError if true_mask is not a single square 2D mask.
      """
      if true_mask.ndim != 2:
          raise ValueError(
              f"Expected a single 2D masks, got shape {true_mask.shape}"
          )
      if true_mask.shape[0] != true_mask.shape[1]:
          raise ValueError(
              f"Expected mask to be square, got shape {true_mask.shape}"
          )
  
      cx_true, cy_true, w_true, h_true = jnp.split(true_box, 4, axis=-1)
      cx_pred, cy_pred, w_pred, h_pred = jnp.split(pred_box, 4, axis=-1)
      w_pred = jnp.maximum(w_pred, 1e-6)
      h_pred = jnp.maximum(h_pred, 1e-6)
  
      # Get top-left corner coordinates:
      x0_true = cx_true - w_true / 2
      y0_true = cy_true - h_true / 2
      x0_pred = cx_pred - w_pred / 2
      y0_pred = cy_pred - h_pred / 2
  
      # Get coordinates of the true box in the reference frame of the pred box:
      x0_true_wrt_pred = (x0_true - x0_pred) / w_pred
      y0_true_wrt_pred = (y0_true - y0_pred) / h_pred
      w_true_wrt_pred = w_true / w_pred
      h_true_wrt_pred = h_true / h_pred
  
      # Scale and translate true masks:
      width = true_mask.shape[0]
      return jax.image.scale_and_translate(
          true_mask,
          shape=true_mask.shape,
          spatial_dims=(0, 1),
          scale=jnp.concatenate((h_true_wrt_pred, w_true_wrt_pred)),
          translation=width
          * jnp.concatenate((y0_true_wrt_pred, x0_true_wrt_pred)),
          method="linear",
      )
  
imports: import abc
  import functools
  from typing import Any, Callable, Dict, Optional, Tuple
  import jax
  import jax.numpy
  import ml_collections
  from scenic.model_lib import matchers
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import box_utils
  from scenic.model_lib.base_models import model_utils
  from scenic.projects.owl_vit import losses
label_desc: Use image node_id 1 for calling a global function with example usage: move_true_to_pred_mask() and returns: jax