class_docstr: Tests for modules in token-learner model.py.
docstr: Tests shapes of TokenLearner Encoder (with and without TokenFuser).
code: def test_encoder(
      self,
      input_shape,
      tokenizer_type,
      temporal_dimensions,
      num_tokens,
      encoder_function,
  ):
      """Tests shapes of TokenLearner Encoder (with and without TokenFuser)."""
      rng = random.PRNGKey(0)
      dummy_input = jnp.ones(input_shape)
      encoder = functools.partial(
          encoder_function,
          num_layers=3,
          mlp_dim=192,
          num_heads=3,
          tokenizer_type=tokenizer_type,
          temporal_dimensions=temporal_dimensions,
          num_tokens=num_tokens,
          tokenlearner_loc=2,
      )
      encoder_vars = encoder().init(rng, dummy_input)
      y = encoder().apply(encoder_vars, dummy_input)
  
      if encoder_function == model.EncoderMod:
          if tokenizer_type == "dynamic":
              expected_shape = (
                  input_shape[0],
                  num_tokens,
                  input_shape[2],
              )
          elif tokenizer_type == "video":
              expected_shape = (
                  input_shape[0],
                  num_tokens * temporal_dimensions,
                  input_shape[2],
              )
          else:
              raise ValueError("Unknown tokenizer type.")
      elif encoder_function == model.EncoderModFuser:
          expected_shape = input_shape
      else:
          raise ValueError("Unknown encoder function.")
  
      self.assertEqual(y.shape, expected_shape)
  
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 2 for calling the TokenLearnerTest obj's underlying member method code with example usage: obj.test_encoder(input_shape, tokenizer_type, temporal_dimensions, num_tokens, encoder_function) without return types