class_docstr: Encoder that supports dropped tokens and object-aware attention.
docstr: Subsample token x according to the keeped index in fg_inds.
  
  Args:
    x: batch x num_tokens x hidden_dim
    fg_inds: batch x num_fg_tokens
    token_scores: batch x num_objects x num_tokens
    drop_scores: bool
  Returns:
    x: batch x num_fg_tokens x hidden_dim
    sampled_token_scores: batch x num_objects x num_fg_tokens
docstr_tok: ["Subsample","token","x","according","to","the","keeped","index","in","fg_inds",".","Args",":","x",":","batch","x","num_tokens","x","hidden_dim","fg_inds",":","batch","x","num_fg_tokens","token_scores",":","batch","x","num_objects","x","num_tokens","drop_scores",":","bool","Returns",":","x",":","batch","x","num_fg_tokens","x","hidden_dim","sampled_token_scores",":","batch","x","num_objects","x","num_fg_tokens"]
code: def _drop_tokens(
      self, x, fg_inds, token_scores=None, drop_scores=False
  ):
      """Subsample token x according to the keeped index in fg_inds.
  
      Args:
        x: batch x num_tokens x hidden_dim
        fg_inds: batch x num_fg_tokens
        token_scores: batch x num_objects x num_tokens
        drop_scores: bool
      Returns:
        x: batch x num_fg_tokens x hidden_dim
        sampled_token_scores: batch x num_objects x num_fg_tokens
      """
      batch, num_tokens, hidden_dim = x.shape
      k = fg_inds.shape[1]
      #  This is identical to torch.gather(x, 1, fg_inds)
      inds = (
          jnp.arange(batch * k).reshape(batch, k) // k * num_tokens
          + fg_inds
      )
      x = x.reshape(batch * num_tokens, hidden_dim)
      x = x[inds.reshape(-1)].reshape(batch, k, hidden_dim)
      sampled_token_scores = None
      if drop_scores:
          num_objs = token_scores.shape[1]
          token_scores = token_scores.reshape(-1)
          expand_fg_inds = jnp.broadcast_to(
              fg_inds[:, None], (batch, num_objs, k)
          )
          inds = jnp.arange(
              batch * num_objs * k
          ) // k * num_tokens + expand_fg_inds.reshape(-1)
          sampled_token_scores = token_scores[inds].reshape(
              batch, num_objs, k
          )
      return x, sampled_token_scores
  
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 3 for calling the CustomEncoder obj's underlying member method code with example usage: obj._drop_tokens(x, fg_inds, token_scores, drop_scores) and returns: x,  sampled_token_scores