class_docstr: A layer that selects top-k tokens.
  
  Note that if `pool_unselected_tokens` is set to True, it pools all the
  unselected tokens and appends it as an extra tokens and returns k+1 tokens.
  
  Attributes:
    top_k: Number of tokens we select.
    sample_tokens: Whether sample the top-k tokens given their scores or just
      take the top-k.
    pool_unselected_tokens: Whether we pool the unselected tokens and attach the
      pooled version as an extra token to the selected tokens.
    exclude_cls: If set to True, it assumes the token at position 0 is CLS token
      and should be excluded from the selection process and be attached back at
      the end.
    score_net_kernel_init: Kernel initialization for the score net.
    dtype: Jax dtype.
class_docstr_tok: ["A","layer","that","selects","top-k","tokens",".","Note","that","if","`","pool_unselected_tokens","`","is","set","to","True",",","it","pools","all","the","unselected","tokens","and","appends","it","as","an","extra","tokens","and","returns","k+1","tokens",".","Attributes",":","top_k",":","Number","of","tokens","we","select",".","sample_tokens",":","Whether","sample","the","top-k","tokens","given","their","scores","or","just","take","the","top-k.","pool_unselected_tokens",":","Whether","we","pool","the","unselected","tokens","and","attach","the","pooled","version","as","an","extra","token","to","the","selected","tokens",".","exclude_cls",":","If","set","to","True",",","it","assumes","the","token","at","position","0","is","CLS","token","and","should","be","excluded","from","the","selection","process","and","be","attached","back","at","the","end",".","score_net_kernel_init",":","Kernel","initialization","for","the","score","net",".","dtype",":","Jax","dtype","."]
code: def __call__(
      self, inputs: jnp.ndarray, *, train: bool
  ) -> jnp.ndarray:
      if self.exclude_cls:
          cls, inputs = jnp.split(inputs, (1,), axis=1)
      input_len = inputs.shape[1]
      if self.top_k > input_len:
          raise ValueError(
              f"The value of top_{self.top_k} should be less than"
              f"input length:{input_len}."
          )
      logging.info(
          "Selecting %d tokens out of %d tokens.", self.top_k, input_len
      )
      # TODO(dehghani): Explore if adding a non-linearity to the score_net helps.
      score_logits = jnp.squeeze(
          nn.Dense(
              features=1,
              dtype=self.dtype,
              kernel_init=self.score_net_kernel_init,
              # No bias is needed since it gets removed during normalization.
              use_bias=False,
              name="score_net",
          )(inputs),
          axis=-1,
      )
  
      if train and self.sample_tokens and self.top_k < input_len:
          # We use dropout rng for sampling, which is always provided.
          rng = self.make_rng("dropout")
          selected_index = sample_categorical(
              rng, score_logits, self.top_k, replacement=False
          )
          selected_logits = jax.vmap(jnp.take, (0, 0, None))(
              score_logits, selected_index, 0
          )
      else:
          selected_logits, selected_index = jax.lax.top_k(
              score_logits, self.top_k
          )
  
      # Take selected tokens:
      selected_tokens = jax.vmap(jnp.take, (0, 0, None))(
          inputs, selected_index, 0
      )
      # Normalize "selected logits" and used as weights for selected tokens:
      selected_tokens = (
          selected_tokens
          * jax.nn.softmax(selected_logits)[..., jnp.newaxis]
      )
  
      if self.pool_unselected_tokens and self.top_k < input_len:
          # Extract index of unselected tokens:
          selected_index_one_hot = jax.nn.one_hot(
              selected_index, num_classes=input_len, dtype=jnp.bool_
          )
          unselected_index_one_hot = jnp.any(
              jnp.logical_not(selected_index_one_hot), axis=1
          )
          _, unselected_index = jax.lax.top_k(
              unselected_index_one_hot, input_len - self.top_k
          )
  
          # Take unselected tokens:
          unselected_tokens = jax.vmap(jnp.take, (0, 0, None))(
              inputs, unselected_index, 0
          )
          unselected_logits = jax.vmap(jnp.take, (0, 0, None))(
              score_logits, unselected_index, 0
          )
          # Normalize "unselected logits" and used as weights for unselected tokens:
          weighted_unselected_tokens = (
              unselected_tokens
              * jax.nn.softmax(unselected_logits)[..., jnp.newaxis]
          )
          unselected_tokens_rep = jnp.sum(
              weighted_unselected_tokens, axis=1, keepdims=True
          )
  
          selected_tokens = jnp.concatenate(
              [selected_tokens, unselected_tokens_rep], axis=1
          )
  
      if self.exclude_cls:
          selected_tokens = jnp.concatenate(
              [cls, selected_tokens], axis=1
          )
      return selected_tokens
  
imports: import abc
  import enum
  import functools
  from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Iterable
  from absl import logging
  from flax import linen
  import jax
  from jax import lax
  from jax import random
  import jax.numpy
  import numpy
  from scenic.model_lib.layers import attention_layers
label_desc: Use image node_id 1 for calling the TopKTokenSelector obj's underlying member method code with example usage: obj.__call__(inputs) and returns: selected_tokens