class_docstr: Tests for token selector module.
class_docstr_tok: ["Tests","for","token","selector","module","."]
docstr: Tests Top-K selector.
docstr_tok: ["Tests","Top-K","selector","."]
code: def test_top_k_selector(
      self,
      pool_unselected_tokens,
      exclude_cls,
      sample_tokens,
      expected_output_len,
  ):
      """Tests Top-K selector."""
      rng, sample_rng = random.split(random.PRNGKey(0))
      x = jnp.ones((4, 16, 32))
      top_k = 5
      top_selector = functools.partial(
          model_utils.TopKTokenSelector,
          top_k=top_k,
          sample_tokens=sample_tokens,
          pool_unselected_tokens=pool_unselected_tokens,
          exclude_cls=exclude_cls,
      )
      variable = top_selector().init(rng, x, train=False)
      y = top_selector().apply(
          variable, x, train=True, rngs={"dropout": sample_rng}
      )
      # Test outputs shape.
      expected_shape = (4, expected_output_len, 32)
      self.assertEqual(y.shape, expected_shape)
  
file_path: scenic/scenic/projects/fast_vit/tests/test_model_utils.py
imports: import functools
  from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  from jax import random
  import jax.numpy
  import numpy
  from scenic.projects.fast_vit import model_utils
label_desc: Use image node_id 1 for calling the TopKTokenSelectorTest obj's underlying member method code with example usage: obj.test_top_k_selector(pool_unselected_tokens, exclude_cls, sample_tokens, expected_output_len) without return types