code: def compress_masks(mask_probs, k=3):
      """At each pixel, stores the largest k probabilities and their indices."""
      if mask_probs.ndim == 5:
          mask_probs = jnp.squeeze(
              mask_probs, axis=-1
          )  # Remove channel dim.
      # Input shape should be [b, num_queries, out_h, out_w].
      assert (
          mask_probs.ndim == 4
      ), f"Expected 4-D input, got {mask_probs.shape}"
      mask_probs = jnp.transpose(mask_probs, [0, 2, 3, 1])
      vals, inds = jax.lax.top_k(mask_probs, k=k)
      # Back to [b, k, out_h, out_w]
      vals = jnp.transpose(vals, [0, 3, 1, 2])
      inds = jnp.transpose(inds, [0, 3, 1, 2])
      return vals, inds
  
imports: from typing import Optional
  import jax
  import jax.numpy
  import numpy
  from PIL import Image
label_desc: Use image node_id 1 for calling a global function with example usage: compress_masks(mask_probs, k) and returns: vals,  inds