code: def get_object_inds(
      scores, num_tokens_per_frame, configs, factorized_encoder=False
  ):
      """Generate inds from scores.
  
      Args:
        scores: batch x num_objs x num_tokens
        num_tokens_per_frame: int
        configs: config dict
        factorized_encoder: bool; if it is used in the factorized_encoder model.
          Get index for each frame if so.
      Returns:
        inds: batch x num_attach_tokens if factorized_encoder==False else
          batch x num_frames x num_frame_attach_tokens
      """
      num_total_attach_tokens = configs.get(
          "num_total_attach_tokens", -1
      )
      num_frame_attach_tokens = configs.get(
          "num_frame_attach_tokens", -1
      )
      assert (
          num_total_attach_tokens == -1 or num_frame_attach_tokens == -1
      )
      assert (not factorized_encoder) or num_frame_attach_tokens > 0
      batch = scores.shape[0]
      # combine heatmaps from different objects
      pooled_scores = scores.max(axis=1)  # batch x num_tokens
      if num_total_attach_tokens > 0:
          pooled_scores = pooled_scores.reshape(batch, -1)
          _, inds = jax.lax.top_k(
              pooled_scores, k=num_total_attach_tokens
          )
      else:
          pooled_scores = pooled_scores.reshape(
              batch, -1, num_tokens_per_frame
          )  # B x T x HW
          num_frames = pooled_scores.shape[1]
          _, inds = jax.lax.top_k(
              pooled_scores, k=num_frame_attach_tokens
          )  # B x T x k
          if not factorized_encoder:
              inds = inds.reshape(batch, -1)  # [B * num_objs, Tk]
              base = (
                  jnp.arange(num_frames * num_frame_attach_tokens)
                  // num_frame_attach_tokens
                  * num_tokens_per_frame
              )
              inds = inds + base[None]  # [B * num_objs, Tk]
              inds = inds.reshape(batch, -1)
      return inds
  
imports: import functools
  from typing import Any, Callable, Iterable, Optional
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  from scenic.model_lib.layers import attention_layers
  from scenic.projects.objectvivit.object_attention import ObjectBlock
label_desc: Use image node_id 1 for calling a global function with example usage: get_object_inds(scores, num_tokens_per_frame, configs, factorized_encoder) and returns: inds