class_docstr: Encoder that supports dropped tokens and object-aware attention.
docstr: Append random/uniformly sampled token from full_x to x.
  
  Args:
    full_x: batch x num_tokens x hidden_dim
    num_context_tokens: number of context tokens to add
    fg_inds: batch x num_fg_tokens
  Returns:
    context_tokens: batch x num_context_token x hidden_dim
docstr_tok: ["Append","random\/uniformly","sampled","token","from","full_x","to","x.","Args",":","full_x",":","batch","x","num_tokens","x","hidden_dim","num_context_tokens",":","number","of","context","tokens","to","add","fg_inds",":","batch","x","num_fg_tokens","Returns",":","context_tokens",":","batch","x","num_context_token","x","hidden_dim"]
code: def _add_context_tokens(
      self, full_x, num_context_tokens, fg_inds=None
  ):
      """Append random/uniformly sampled token from full_x to x.
  
      Args:
        full_x: batch x num_tokens x hidden_dim
        num_context_tokens: number of context tokens to add
        fg_inds: batch x num_fg_tokens
      Returns:
        context_tokens: batch x num_context_token x hidden_dim
      """
      batch, num_total_tokens, hidden_dim = full_x.shape
      k = fg_inds.shape[1]
      random_score = jax.random.uniform(
          self.make_rng("dropout"), (batch, num_total_tokens)
      )
  
      obj_inds = (
          jnp.arange(batch * k).reshape(batch, k)
          // k
          * num_total_tokens
          + fg_inds
      )
      random_score = random_score.reshape(-1)
      random_score = random_score.at[obj_inds].set(0)
      random_score = random_score.reshape(batch, num_total_tokens)
  
      _, random_inds = jax.lax.top_k(random_score, k=num_context_tokens)
      base = (
          jnp.arange(batch * num_context_tokens).reshape(
              batch, num_context_tokens
          )
          // num_context_tokens
          * num_total_tokens
      )
      inds = base + random_inds
      full_x = full_x.reshape(batch * num_total_tokens, hidden_dim)
      context_tokens = full_x[inds.reshape(-1)]
      context_tokens = context_tokens.reshape(
          batch, num_context_tokens, hidden_dim
      )
      return context_tokens
  
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 2 for calling the CustomEncoder obj's underlying member method code with example usage: obj._add_context_tokens(full_x, num_context_tokens, fg_inds) and returns: context_tokens