class_docstr: Tests for modules in token-learner model.py.
docstr: Tests TokenLearner module.
code: def test_dynamic_tokenizer(self, num_tokens):
      """Tests TokenLearner module."""
      rng = random.PRNGKey(0)
      x = jnp.ones((4, 224, 224, 64))
      tokenizer = functools.partial(
          model.TokenLearnerModule, num_tokens=num_tokens
      )
      tokenizer_vars = tokenizer().init(rng, x)
      y = tokenizer().apply(tokenizer_vars, x)
      # Test outputs shape.
      self.assertEqual(y.shape, (x.shape[0], num_tokens, x.shape[-1]))
  
file_path: scenic/scenic/projects/token_learner/tests/test_model.py
imports: import functools
  from absl.testing import absltest
  from absl.testing import parameterized
  from jax import random
  import jax.numpy
  from scenic.projects.token_learner import model
label_desc: Use image node_id 1 for calling the TokenLearnerTest obj's underlying member method code with example usage: obj.test_dynamic_tokenizer(num_tokens) without return types